mrleo1nid/ACE-Step
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

work on pip package

Browse Source
This commit is contained in:
mrfakename
2025-05-06 18:59:32 -07:00
parent 54da683d36
commit a5746eaab6
36 changed files with 2928 additions and 1251 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/models/__init__.py
View File
acestep/models/ace_step_transformer.py
+155 -46
View File
@@ -31,9 +31,15 @@ from .lyrics_utils.lyric_encoder import ConformerEncoder as LyricEncoder
def cross_norm(hidden_states, controlnet_input):
# input N x T x c
mean_hidden_states, std_hidden_states = hidden_states.mean(dim=(1,2), keepdim=True), hidden_states.std(dim=(1,2), keepdim=True)
mean_controlnet_input, std_controlnet_input = controlnet_input.mean(dim=(1,2), keepdim=True), controlnet_input.std(dim=(1,2), keepdim=True)
controlnet_input = (controlnet_input - mean_controlnet_input) * (std_hidden_states / (std_controlnet_input + 1e-12)) + mean_hidden_states
mean_hidden_states, std_hidden_states = hidden_states.mean(
dim=(1, 2), keepdim=True
), hidden_states.std(dim=(1, 2), keepdim=True)
mean_controlnet_input, std_controlnet_input = controlnet_input.mean(
dim=(1, 2), keepdim=True
), controlnet_input.std(dim=(1, 2), keepdim=True)
controlnet_input = (controlnet_input - mean_controlnet_input) * (
std_hidden_states / (std_controlnet_input + 1e-12)
) + mean_hidden_states
return controlnet_input
@@ -45,17 +51,27 @@ class Qwen2RotaryEmbedding(nn.Module):
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.base = base
inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
inv_freq = 1.0 / (
self.base
** (
torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device)
/ self.dim
)
)
self.register_buffer("inv_freq", inv_freq, persistent=False)
# Build here to make `torch.jit.trace` work.
self._set_cos_sin_cache(
seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
seq_len=max_position_embeddings,
device=self.inv_freq.device,
dtype=torch.get_default_dtype(),
)
def _set_cos_sin_cache(self, seq_len, device, dtype):
self.max_seq_len_cached = seq_len
t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
t = torch.arange(
self.max_seq_len_cached, device=device, dtype=torch.int64
).type_as(self.inv_freq)
freqs = torch.outer(t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
@@ -82,8 +98,12 @@ class T2IFinalLayer(nn.Module):
def __init__(self, hidden_size, patch_size=[16, 1], out_channels=256):
super().__init__()
self.norm_final = nn.RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.linear = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True)
self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5)
self.linear = nn.Linear(
hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True
)
self.scale_shift_table = nn.Parameter(
torch.randn(2, hidden_size) / hidden_size**0.5
)
self.out_channels = out_channels
self.patch_size = patch_size
@@ -95,14 +115,28 @@ class T2IFinalLayer(nn.Module):
# 4 unpatchify
new_height, new_width = 1, hidden_states.size(1)
hidden_states = hidden_states.reshape(
shape=(hidden_states.shape[0], new_height, new_width, self.patch_size[0], self.patch_size[1], self.out_channels)
shape=(
hidden_states.shape[0],
new_height,
new_width,
self.patch_size[0],
self.patch_size[1],
self.out_channels,
)
).contiguous()
hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
output = hidden_states.reshape(
shape=(hidden_states.shape[0], self.out_channels, new_height * self.patch_size[0], new_width * self.patch_size[1])
shape=(
hidden_states.shape[0],
self.out_channels,
new_height * self.patch_size[0],
new_width * self.patch_size[1],
)
).contiguous()
if width > new_width:
output = torch.nn.functional.pad(output, (0, width - new_width, 0, 0), 'constant', 0)
output = torch.nn.functional.pad(
output, (0, width - new_width, 0, 0), "constant", 0
)
elif width < new_width:
output = output[:, :, :, :width]
return output
@@ -131,9 +165,25 @@ class PatchEmbed(nn.Module):
super().__init__()
patch_size_h, patch_size_w = patch_size
self.early_conv_layers = nn.Sequential(
nn.Conv2d(in_channels, in_channels*256, kernel_size=patch_size, stride=patch_size, padding=0, bias=bias),
torch.nn.GroupNorm(num_groups=32, num_channels=in_channels*256, eps=1e-6, affine=True),
nn.Conv2d(in_channels*256, embed_dim, kernel_size=1, stride=1, padding=0, bias=bias)
nn.Conv2d(
in_channels,
in_channels * 256,
kernel_size=patch_size,
stride=patch_size,
padding=0,
bias=bias,
),
torch.nn.GroupNorm(
num_groups=32, num_channels=in_channels * 256, eps=1e-6, affine=True
),
nn.Conv2d(
in_channels * 256,
embed_dim,
kernel_size=1,
stride=1,
padding=0,
bias=bias,
),
)
self.patch_size = patch_size
self.height, self.width = height // patch_size_h, width // patch_size_w
@@ -153,7 +203,9 @@ class Transformer2DModelOutput(BaseOutput):
proj_losses: Optional[Tuple[Tuple[str, torch.Tensor]]] = None
class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
class ACEStepTransformer2DModel(
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
):
_supports_gradient_checkpointing = True
@register_to_config
@@ -217,9 +269,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
)
self.num_layers = num_layers
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim)
self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True))
self.time_proj = Timesteps(
num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0
)
self.timestep_embedder = TimestepEmbedding(
in_channels=256, time_embed_dim=self.inner_dim
)
self.t_block = nn.Sequential(
nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True)
)
# speaker
self.speaker_embedder = nn.Linear(speaker_embedding_dim, self.inner_dim)
@@ -229,25 +287,30 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# lyric
self.lyric_embs = nn.Embedding(lyric_encoder_vocab_size, lyric_hidden_size)
self.lyric_encoder = LyricEncoder(input_size=lyric_hidden_size, static_chunk_size=0)
self.lyric_encoder = LyricEncoder(
input_size=lyric_hidden_size, static_chunk_size=0
)
self.lyric_proj = nn.Linear(lyric_hidden_size, self.inner_dim)
projector_dim = 2 * self.inner_dim
self.projectors = nn.ModuleList([
nn.Sequential(
nn.Linear(self.inner_dim, projector_dim),
nn.SiLU(),
nn.Linear(projector_dim, projector_dim),
nn.SiLU(),
nn.Linear(projector_dim, ssl_dim),
) for ssl_dim in ssl_latent_dims
])
self.projectors = nn.ModuleList(
[
nn.Sequential(
nn.Linear(self.inner_dim, projector_dim),
nn.SiLU(),
nn.Linear(projector_dim, projector_dim),
nn.SiLU(),
nn.Linear(projector_dim, ssl_dim),
)
for ssl_dim in ssl_latent_dims
]
)
self.ssl_latent_dims = ssl_latent_dims
self.ssl_encoder_depths = ssl_encoder_depths
self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction='mean')
self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction="mean")
self.ssl_names = ssl_names
self.proj_in = PatchEmbed(
@@ -258,11 +321,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
bias=True,
)
self.final_layer = T2IFinalLayer(self.inner_dim, patch_size=patch_size, out_channels=out_channels)
self.final_layer = T2IFinalLayer(
self.inner_dim, patch_size=patch_size, out_channels=out_channels
)
self.gradient_checkpointing = False
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
def enable_forward_chunking(
self, chunk_size: Optional[int] = None, dim: int = 0
) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -281,7 +348,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# By default chunk size is 1
chunk_size = chunk_size or 1
def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
def fn_recursive_feed_forward(
module: torch.nn.Module, chunk_size: int, dim: int
):
if hasattr(module, "set_chunk_feed_forward"):
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
@@ -302,7 +371,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
):
# N x T x D
lyric_embs = self.lyric_embs(lyric_token_idx)
prompt_prenet_out, _mask = self.lyric_encoder(lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1)
prompt_prenet_out, _mask = self.lyric_encoder(
lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1
)
prompt_prenet_out = self.lyric_proj(prompt_prenet_out)
return prompt_prenet_out
@@ -317,7 +388,7 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
bs = encoder_text_hidden_states.shape[0]
device = encoder_text_hidden_states.device
# speaker embedding
encoder_spk_hidden_states = self.speaker_embedder(speaker_embeds).unsqueeze(1)
speaker_mask = torch.ones(bs, 1, device=device)
@@ -331,8 +402,17 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
lyric_mask=lyric_mask,
)
encoder_hidden_states = torch.cat([encoder_spk_hidden_states, encoder_text_hidden_states, encoder_lyric_hidden_states], dim=1)
encoder_hidden_mask = torch.cat([speaker_mask, text_attention_mask, lyric_mask], dim=1)
encoder_hidden_states = torch.cat(
[
encoder_spk_hidden_states,
encoder_text_hidden_states,
encoder_lyric_hidden_states,
],
dim=1,
)
encoder_hidden_mask = torch.cat(
[speaker_mask, text_attention_mask, lyric_mask], dim=1
)
return encoder_hidden_states, encoder_hidden_mask
def decode(
@@ -344,12 +424,16 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
timestep: Optional[torch.Tensor],
ssl_hidden_states: Optional[List[torch.Tensor]] = None,
output_length: int = 0,
block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
block_controlnet_hidden_states: Optional[
Union[List[torch.Tensor], torch.Tensor]
] = None,
controlnet_scale: Union[float, torch.Tensor] = 1.0,
return_dict: bool = True,
):
embedded_timestep = self.timestep_embedder(self.time_proj(timestep).to(dtype=hidden_states.dtype))
embedded_timestep = self.timestep_embedder(
self.time_proj(timestep).to(dtype=hidden_states.dtype)
)
temb = self.t_block(embedded_timestep)
hidden_states = self.proj_in(hidden_states)
@@ -361,8 +445,12 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
inner_hidden_states = []
rotary_freqs_cis = self.rotary_emb(hidden_states, seq_len=hidden_states.shape[1])
encoder_rotary_freqs_cis = self.rotary_emb(encoder_hidden_states, seq_len=encoder_hidden_states.shape[1])
rotary_freqs_cis = self.rotary_emb(
hidden_states, seq_len=hidden_states.shape[1]
)
encoder_rotary_freqs_cis = self.rotary_emb(
encoder_hidden_states, seq_len=encoder_hidden_states.shape[1]
)
for index_block, block in enumerate(self.transformer_blocks):
@@ -377,7 +465,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
return custom_forward
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
ckpt_kwargs: Dict[str, Any] = (
{"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
)
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states=hidden_states,
@@ -406,9 +496,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
inner_hidden_states.append(hidden_states)
proj_losses = []
if len(inner_hidden_states) > 0 and ssl_hidden_states is not None and len(ssl_hidden_states) > 0:
if (
len(inner_hidden_states) > 0
and ssl_hidden_states is not None
and len(ssl_hidden_states) > 0
):
for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names):
for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(
inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names
):
if ssl_hidden_state is None:
continue
# 1. N x T x D1 -> N x D x D2
@@ -416,9 +512,20 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# 3. projection loss
bs = inner_hidden_state.shape[0]
proj_loss = 0.0
for i, (z, z_tilde) in enumerate(zip(ssl_hidden_state, est_ssl_hidden_state)):
for i, (z, z_tilde) in enumerate(
zip(ssl_hidden_state, est_ssl_hidden_state)
):
# 2. interpolate
z_tilde = F.interpolate(z_tilde.unsqueeze(0).transpose(1, 2), size=len(z), mode='linear', align_corners=False).transpose(1, 2).squeeze(0)
z_tilde = (
F.interpolate(
z_tilde.unsqueeze(0).transpose(1, 2),
size=len(z),
mode="linear",
align_corners=False,
)
.transpose(1, 2)
.squeeze(0)
)
z_tilde = torch.nn.functional.normalize(z_tilde, dim=-1)
z = torch.nn.functional.normalize(z, dim=-1)
@@ -445,7 +552,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
lyric_mask: Optional[torch.LongTensor] = None,
timestep: Optional[torch.Tensor] = None,
ssl_hidden_states: Optional[List[torch.Tensor]] = None,
block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
block_controlnet_hidden_states: Optional[
Union[List[torch.Tensor], torch.Tensor]
] = None,
controlnet_scale: Union[float, torch.Tensor] = 1.0,
return_dict: bool = True,
):
acestep/models/attention.py
+15 -3
View File
@@ -23,10 +23,18 @@ from diffusers.models.normalization import RMSNorm
try:
# from .dcformer import DCMHAttention
from .customer_attention_processor import Attention, CustomLiteLAProcessor2_0, CustomerAttnProcessor2_0
from .customer_attention_processor import (
Attention,
CustomLiteLAProcessor2_0,
CustomerAttnProcessor2_0,
)
except ImportError:
# from dcformer import DCMHAttention
from customer_attention_processor import Attention, CustomLiteLAProcessor2_0, CustomerAttnProcessor2_0
from customer_attention_processor import (
Attention,
CustomLiteLAProcessor2_0,
CustomerAttnProcessor2_0,
)
logger = logging.get_logger(__name__)
@@ -55,13 +63,16 @@ def t2i_modulate(x, shift, scale):
return x * (1 + scale) + shift
def get_same_padding(kernel_size: Union[int, Tuple[int, ...]]) -> Union[int, Tuple[int, ...]]:
def get_same_padding(
kernel_size: Union[int, Tuple[int, ...]],
) -> Union[int, Tuple[int, ...]]:
if isinstance(kernel_size, tuple):
return tuple([get_same_padding(ks) for ks in kernel_size])
else:
assert kernel_size % 2 > 0, f"kernel size {kernel_size} should be odd number"
return kernel_size // 2
class ConvLayer(nn.Module):
def __init__(
self,
@@ -187,6 +198,7 @@ class LinearTransformerBlock(nn.Module):
"""
A Sana block with global shared adaptive layer norm (adaLN-single) conditioning.
"""
def __init__(
self,
dim,
acestep/models/customer_attention_processor.py
+106 -31
View File
@@ -78,12 +78,16 @@ class CustomLiteLAProcessor2_0:
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
hidden_states = hidden_states.view(
batch_size, channel, height * width
).transpose(1, 2)
if encoder_hidden_states is not None:
context_input_ndim = encoder_hidden_states.ndim
if context_input_ndim == 4:
batch_size, channel, height, width = encoder_hidden_states.shape
encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
encoder_hidden_states = encoder_hidden_states.view(
batch_size, channel, height * width
).transpose(1, 2)
batch_size = hidden_states.shape[0]
@@ -94,7 +98,11 @@ class CustomLiteLAProcessor2_0:
value = attn.to_v(hidden_states)
# `context` projections.
has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
has_encoder_hidden_state_proj = (
hasattr(attn, "add_q_proj")
and hasattr(attn, "add_k_proj")
and hasattr(attn, "add_v_proj")
)
if encoder_hidden_states is not None and has_encoder_hidden_state_proj:
encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
@@ -114,7 +122,11 @@ class CustomLiteLAProcessor2_0:
head_dim = inner_dim // attn.heads
query = query.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
key = key.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1).transpose(-1, -2)
key = (
key.transpose(-1, -2)
.reshape(batch_size, attn.heads, head_dim, -1)
.transpose(-1, -2)
)
value = value.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
# RoPE需要 [B, H, S, D] 输入
@@ -140,17 +152,33 @@ class CustomLiteLAProcessor2_0:
if attention_mask is not None:
# attention_mask: [B, S] -> [B, 1, S, 1]
attention_mask = attention_mask[:, None, :, None].to(key.dtype) # [B, 1, S, 1]
query = query * attention_mask.permute(0, 1, 3, 2) # [B, H, S, D] * [B, 1, S, 1]
attention_mask = attention_mask[:, None, :, None].to(
key.dtype
) # [B, 1, S, 1]
query = query * attention_mask.permute(
0, 1, 3, 2
) # [B, H, S, D] * [B, 1, S, 1]
if not attn.is_cross_attention:
key = key * attention_mask # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘
value = value * attention_mask.permute(0, 1, 3, 2) # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度
key = (
key * attention_mask
) # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘
value = value * attention_mask.permute(
0, 1, 3, 2
) # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度
if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(key.dtype) # [B, 1, S_enc, 1]
if (
attn.is_cross_attention
and encoder_attention_mask is not None
and has_encoder_hidden_state_proj
):
encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(
key.dtype
) # [B, 1, S_enc, 1]
# 此时 key: [B, h, S_enc, D], value: [B, h, D, S_enc]
key = key * encoder_attention_mask # [B, h, S_enc, D] * [B, 1, S_enc, 1]
value = value * encoder_attention_mask.permute(0, 1, 3, 2) # [B, h, D, S_enc] * [B, 1, 1, S_enc]
value = value * encoder_attention_mask.permute(
0, 1, 3, 2
) # [B, h, D, S_enc] * [B, 1, 1, S_enc]
query = self.kernel_func(query)
key = self.kernel_func(key)
@@ -168,16 +196,22 @@ class CustomLiteLAProcessor2_0:
hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps)
hidden_states = hidden_states.view(batch_size, attn.heads * head_dim, -1).permute(0, 2, 1)
hidden_states = hidden_states.view(
batch_size, attn.heads * head_dim, -1
).permute(0, 2, 1)
hidden_states = hidden_states.to(dtype)
if encoder_hidden_states is not None:
encoder_hidden_states = encoder_hidden_states.to(dtype)
# Split the attention outputs.
if encoder_hidden_states is not None and not attn.is_cross_attention and has_encoder_hidden_state_proj:
if (
encoder_hidden_states is not None
and not attn.is_cross_attention
and has_encoder_hidden_state_proj
):
hidden_states, encoder_hidden_states = (
hidden_states[:, : hidden_states_len],
hidden_states[:, :hidden_states_len],
hidden_states[:, hidden_states_len:],
)
@@ -185,13 +219,22 @@ class CustomLiteLAProcessor2_0:
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if encoder_hidden_states is not None and not attn.context_pre_only and not attn.is_cross_attention and hasattr(attn, "to_add_out"):
if (
encoder_hidden_states is not None
and not attn.context_pre_only
and not attn.is_cross_attention
and hasattr(attn, "to_add_out")
):
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
hidden_states = hidden_states.transpose(-1, -2).reshape(
batch_size, channel, height, width
)
if encoder_hidden_states is not None and context_input_ndim == 4:
encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(
batch_size, channel, height, width
)
if torch.get_autocast_gpu_dtype() == torch.float16:
hidden_states = hidden_states.clip(-65504, 65504)
@@ -208,7 +251,9 @@ class CustomerAttnProcessor2_0:
def __init__(self):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
raise ImportError(
"AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
)
def apply_rotary_emb(
self,
@@ -258,23 +303,35 @@ class CustomerAttnProcessor2_0:
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
hidden_states = hidden_states.view(
batch_size, channel, height * width
).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
hidden_states.shape
if encoder_hidden_states is None
else encoder_hidden_states.shape
)
has_encoder_hidden_state_proj = (
hasattr(attn, "add_q_proj")
and hasattr(attn, "add_k_proj")
and hasattr(attn, "add_v_proj")
)
has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
if attn.group_norm is not None:
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
1, 2
)
query = attn.to_q(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
elif attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
encoder_hidden_states = attn.norm_encoder_hidden_states(
encoder_hidden_states
)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
@@ -300,19 +357,33 @@ class CustomerAttnProcessor2_0:
elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:
key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)
if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
if (
attn.is_cross_attention
and encoder_attention_mask is not None
and has_encoder_hidden_state_proj
):
# attention_mask: N x S1
# encoder_attention_mask: N x S2
# cross attention 整合attention_mask和encoder_attention_mask
combined_mask = attention_mask[:, :, None] * encoder_attention_mask[:, None, :]
combined_mask = (
attention_mask[:, :, None] * encoder_attention_mask[:, None, :]
)
attention_mask = torch.where(combined_mask == 1, 0.0, -torch.inf)
attention_mask = attention_mask[:, None, :, :].expand(-1, attn.heads, -1, -1).to(query.dtype)
attention_mask = (
attention_mask[:, None, :, :]
.expand(-1, attn.heads, -1, -1)
.to(query.dtype)
)
elif not attn.is_cross_attention and attention_mask is not None:
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
attention_mask = attn.prepare_attention_mask(
attention_mask, sequence_length, batch_size
)
# scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
attention_mask = attention_mask.view(
batch_size, attn.heads, -1, attention_mask.shape[-1]
)
# the output of sdp = (batch, num_heads, seq_len, head_dim)
# TODO: add support for attn.scale when we move to Torch 2.1
@@ -320,7 +391,9 @@ class CustomerAttnProcessor2_0:
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.transpose(1, 2).reshape(
batch_size, -1, attn.heads * head_dim
)
hidden_states = hidden_states.to(query.dtype)
# linear proj
@@ -329,7 +402,9 @@ class CustomerAttnProcessor2_0:
hidden_states = attn.to_out[1](hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
hidden_states = hidden_states.transpose(-1, -2).reshape(
batch_size, channel, height, width
)
if attn.residual_connection:
hidden_states = hidden_states + residual
acestep/models/lyrics_utils/__init__.py
View File
acestep/models/lyrics_utils/lyric_encoder.py
+167 -147
View File
@@ -3,16 +3,19 @@ import math
import torch
from torch import nn
class ConvolutionModule(nn.Module):
"""ConvolutionModule in Conformer model."""
def __init__(self,
channels: int,
kernel_size: int = 15,
activation: nn.Module = nn.ReLU(),
norm: str = "batch_norm",
causal: bool = False,
bias: bool = True):
def __init__(
self,
channels: int,
kernel_size: int = 15,
activation: nn.Module = nn.ReLU(),
norm: str = "batch_norm",
causal: bool = False,
bias: bool = True,
):
"""Construct an ConvolutionModule object.
Args:
channels (int): The number of channels of conv layers.
@@ -51,7 +54,7 @@ class ConvolutionModule(nn.Module):
bias=bias,
)
assert norm in ['batch_norm', 'layer_norm']
assert norm in ["batch_norm", "layer_norm"]
if norm == "batch_norm":
self.use_layer_norm = False
self.norm = nn.BatchNorm1d(channels)
@@ -95,13 +98,13 @@ class ConvolutionModule(nn.Module):
if self.lorder > 0:
if cache.size(2) == 0: # cache_t == 0
x = nn.functional.pad(x, (self.lorder, 0), 'constant', 0.0)
x = nn.functional.pad(x, (self.lorder, 0), "constant", 0.0)
else:
assert cache.size(0) == x.size(0) # equal batch
assert cache.size(1) == x.size(1) # equal channel
x = torch.cat((cache, x), dim=2)
assert (x.size(2) > self.lorder)
new_cache = x[:, :, -self.lorder:]
assert x.size(2) > self.lorder
new_cache = x[:, :, -self.lorder :]
else:
# It's better we just return None if no cache is required,
# However, for JIT export, here we just fake one tensor instead of
@@ -126,6 +129,7 @@ class ConvolutionModule(nn.Module):
return x.transpose(1, 2), new_cache
class PositionwiseFeedForward(torch.nn.Module):
"""Positionwise feed forward layer.
@@ -140,11 +144,11 @@ class PositionwiseFeedForward(torch.nn.Module):
"""
def __init__(
self,
idim: int,
hidden_units: int,
dropout_rate: float,
activation: torch.nn.Module = torch.nn.ReLU(),
self,
idim: int,
hidden_units: int,
dropout_rate: float,
activation: torch.nn.Module = torch.nn.ReLU(),
):
"""Construct a PositionwiseFeedForward object."""
super(PositionwiseFeedForward, self).__init__()
@@ -163,6 +167,7 @@ class PositionwiseFeedForward(torch.nn.Module):
"""
return self.w_2(self.dropout(self.activation(self.w_1(xs))))
class Swish(torch.nn.Module):
"""Construct an Swish object."""
@@ -170,6 +175,7 @@ class Swish(torch.nn.Module):
"""Return Swish activation function."""
return x * torch.sigmoid(x)
class MultiHeadedAttention(nn.Module):
"""Multi-Head Attention layer.
@@ -180,11 +186,9 @@ class MultiHeadedAttention(nn.Module):
"""
def __init__(self,
n_head: int,
n_feat: int,
dropout_rate: float,
key_bias: bool = True):
def __init__(
self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True
):
"""Construct an MultiHeadedAttention object."""
super().__init__()
assert n_feat % n_head == 0
@@ -229,7 +233,7 @@ class MultiHeadedAttention(nn.Module):
self,
value: torch.Tensor,
scores: torch.Tensor,
mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
) -> torch.Tensor:
"""Compute attention context vector.
@@ -247,23 +251,24 @@ class MultiHeadedAttention(nn.Module):
"""
n_batch = value.size(0)
if mask.size(2) > 0: # time2 > 0
mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2)
# For last chunk, time2 might be larger than scores.size(-1)
mask = mask[:, :, :, :scores.size(-1)] # (batch, 1, *, time2)
scores = scores.masked_fill(mask, -float('inf'))
mask = mask[:, :, :, : scores.size(-1)] # (batch, 1, *, time2)
scores = scores.masked_fill(mask, -float("inf"))
attn = torch.softmax(scores, dim=-1).masked_fill(
mask, 0.0) # (batch, head, time1, time2)
mask, 0.0
) # (batch, head, time1, time2)
else:
attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2)
p_attn = self.dropout(attn)
x = torch.matmul(p_attn, value) # (batch, head, time1, d_k)
x = (x.transpose(1, 2).contiguous().view(n_batch, -1,
self.h * self.d_k)
) # (batch, time1, d_model)
x = (
x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
) # (batch, time1, d_model)
return self.linear_out(x) # (batch, time1, d_model)
@@ -274,7 +279,7 @@ class MultiHeadedAttention(nn.Module):
value: torch.Tensor,
mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
pos_emb: torch.Tensor = torch.empty(0),
cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Compute scaled dot product attention.
@@ -308,9 +313,7 @@ class MultiHeadedAttention(nn.Module):
"""
q, k, v = self.forward_qkv(query, key, value)
if cache.size(0) > 0:
key_cache, value_cache = torch.split(cache,
cache.size(-1) // 2,
dim=-1)
key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
k = torch.cat([key_cache, k], dim=2)
v = torch.cat([value_cache, v], dim=2)
new_cache = torch.cat((k, v), dim=-1)
@@ -328,11 +331,9 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
dropout_rate (float): Dropout rate.
"""
def __init__(self,
n_head: int,
n_feat: int,
dropout_rate: float,
key_bias: bool = True):
def __init__(
self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True
):
"""Construct an RelPositionMultiHeadedAttention object."""
super().__init__(n_head, n_feat, dropout_rate, key_bias)
# linear transformation for positional encoding
@@ -355,14 +356,12 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
torch.Tensor: Output tensor.
"""
zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1),
device=x.device,
dtype=x.dtype)
zero_pad = torch.zeros(
(x.size()[0], x.size()[1], x.size()[2], 1), device=x.device, dtype=x.dtype
)
x_padded = torch.cat([zero_pad, x], dim=-1)
x_padded = x_padded.view(x.size()[0],
x.size()[1],
x.size(3) + 1, x.size(2))
x_padded = x_padded.view(x.size()[0], x.size()[1], x.size(3) + 1, x.size(2))
x = x_padded[:, :, 1:].view_as(x)[
:, :, :, : x.size(-1) // 2 + 1
] # only keep the positions from 0 to time2
@@ -375,7 +374,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
value: torch.Tensor,
mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
pos_emb: torch.Tensor = torch.empty(0),
cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Compute 'Scaled Dot Product Attention' with rel. positional encoding.
Args:
@@ -399,9 +398,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
q = q.transpose(1, 2) # (batch, time1, head, d_k)
if cache.size(0) > 0:
key_cache, value_cache = torch.split(cache,
cache.size(-1) // 2,
dim=-1)
key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
k = torch.cat([key_cache, k], dim=2)
v = torch.cat([value_cache, v], dim=2)
# NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
@@ -431,15 +428,15 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
matrix_bd = self.rel_shift(matrix_bd)
scores = (matrix_ac + matrix_bd) / math.sqrt(
self.d_k) # (batch, head, time1, time2)
self.d_k
) # (batch, head, time1, time2)
return self.forward_attention(v, scores, mask), new_cache
def subsequent_mask(
size: int,
device: torch.device = torch.device("cpu"),
size: int,
device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
"""Create mask for subsequent steps (size, size).
@@ -474,11 +471,11 @@ def subsequent_mask(
def subsequent_chunk_mask(
size: int,
chunk_size: int,
num_left_chunks: int = -1,
device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
size: int,
chunk_size: int,
num_left_chunks: int = -1,
device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
"""Create mask for subsequent steps (size, size) with chunk size,
this is for streaming encoder
@@ -510,15 +507,18 @@ def subsequent_chunk_mask(
ret[i, start:ending] = True
return ret
def add_optional_chunk_mask(xs: torch.Tensor,
masks: torch.Tensor,
use_dynamic_chunk: bool,
use_dynamic_left_chunk: bool,
decoding_chunk_size: int,
static_chunk_size: int,
num_decoding_left_chunks: int,
enable_full_context: bool = True):
""" Apply optional mask for encoder.
def add_optional_chunk_mask(
xs: torch.Tensor,
masks: torch.Tensor,
use_dynamic_chunk: bool,
use_dynamic_left_chunk: bool,
decoding_chunk_size: int,
static_chunk_size: int,
num_decoding_left_chunks: int,
enable_full_context: bool = True,
):
"""Apply optional mask for encoder.
Args:
xs (torch.Tensor): padded input, (B, L, D), L for max length
@@ -557,7 +557,7 @@ def add_optional_chunk_mask(xs: torch.Tensor,
# chunk size is either [1, 25] or full context(max_len).
# Since we use 4 times subsampling and allow up to 1s(100 frames)
# delay, the maximum frame is 100 / 4 = 25.
chunk_size = torch.randint(1, max_len, (1, )).item()
chunk_size = torch.randint(1, max_len, (1,)).item()
num_left_chunks = -1
if chunk_size > max_len // 2 and enable_full_context:
chunk_size = max_len
@@ -565,18 +565,17 @@ def add_optional_chunk_mask(xs: torch.Tensor,
chunk_size = chunk_size % 25 + 1
if use_dynamic_left_chunk:
max_left_chunks = (max_len - 1) // chunk_size
num_left_chunks = torch.randint(0, max_left_chunks,
(1, )).item()
chunk_masks = subsequent_chunk_mask(xs.size(1), chunk_size,
num_left_chunks,
xs.device) # (L, L)
num_left_chunks = torch.randint(0, max_left_chunks, (1,)).item()
chunk_masks = subsequent_chunk_mask(
xs.size(1), chunk_size, num_left_chunks, xs.device
) # (L, L)
chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L)
chunk_masks = masks & chunk_masks # (B, L, L)
elif static_chunk_size > 0:
num_left_chunks = num_decoding_left_chunks
chunk_masks = subsequent_chunk_mask(xs.size(1), static_chunk_size,
num_left_chunks,
xs.device) # (L, L)
chunk_masks = subsequent_chunk_mask(
xs.size(1), static_chunk_size, num_left_chunks, xs.device
) # (L, L)
chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L)
chunk_masks = masks & chunk_masks # (B, L, L)
else:
@@ -630,7 +629,8 @@ class ConformerEncoderLayer(nn.Module):
if self.conv_module is not None:
self.norm_conv = nn.LayerNorm(size, eps=1e-5) # for the CNN module
self.norm_final = nn.LayerNorm(
size, eps=1e-5) # for the final output of the block
size, eps=1e-5
) # for the final output of the block
self.dropout = nn.Dropout(dropout_rate)
self.size = size
self.normalize_before = normalize_before
@@ -671,8 +671,7 @@ class ConformerEncoderLayer(nn.Module):
residual = x
if self.normalize_before:
x = self.norm_ff_macaron(x)
x = residual + self.ff_scale * self.dropout(
self.feed_forward_macaron(x))
x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x))
if not self.normalize_before:
x = self.norm_ff_macaron(x)
@@ -680,8 +679,7 @@ class ConformerEncoderLayer(nn.Module):
residual = x
if self.normalize_before:
x = self.norm_mha(x)
x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb,
att_cache)
x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, att_cache)
x = residual + self.dropout(x_att)
if not self.normalize_before:
x = self.norm_mha(x)
@@ -712,7 +710,6 @@ class ConformerEncoderLayer(nn.Module):
x = self.norm_final(x)
return x, mask, new_att_cache, new_cnn_cache
class EspnetRelPositionalEncoding(torch.nn.Module):
@@ -770,8 +767,9 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
pe = torch.cat([pe_positive, pe_negative], dim=1)
self.pe = pe.to(device=x.device, dtype=x.dtype)
def forward(self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0) \
-> Tuple[torch.Tensor, torch.Tensor]:
def forward(
self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Add positional encoding.
Args:
@@ -785,11 +783,11 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
x = x * self.xscale
pos_emb = self.position_encoding(size=x.size(1), offset=offset)
return self.dropout(x), self.dropout(pos_emb)
def position_encoding(self,
offset: Union[int, torch.Tensor],
size: int) -> torch.Tensor:
""" For getting encoding in a streaming fashion
def position_encoding(
self, offset: Union[int, torch.Tensor], size: int
) -> torch.Tensor:
"""For getting encoding in a streaming fashion
Attention!!!!!
we apply dropout only once at the whole utterance level in a none
@@ -806,12 +804,11 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
"""
pos_emb = self.pe[
:,
self.pe.size(1) // 2 - size + 1: self.pe.size(1) // 2 + size,
self.pe.size(1) // 2 - size + 1 : self.pe.size(1) // 2 + size,
]
return pos_emb
class LinearEmbed(torch.nn.Module):
"""Linear transform the input without subsampling
@@ -822,8 +819,9 @@ class LinearEmbed(torch.nn.Module):
"""
def __init__(self, idim: int, odim: int, dropout_rate: float,
pos_enc_class: torch.nn.Module):
def __init__(
self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module
):
"""Construct an linear object."""
super().__init__()
self.out = torch.nn.Sequential(
@@ -831,16 +829,15 @@ class LinearEmbed(torch.nn.Module):
torch.nn.LayerNorm(odim, eps=1e-5),
torch.nn.Dropout(dropout_rate),
)
self.pos_enc = pos_enc_class #rel_pos_espnet
def position_encoding(self, offset: Union[int, torch.Tensor],
size: int) -> torch.Tensor:
self.pos_enc = pos_enc_class # rel_pos_espnet
def position_encoding(
self, offset: Union[int, torch.Tensor], size: int
) -> torch.Tensor:
return self.pos_enc.position_encoding(offset, size)
def forward(
self,
x: torch.Tensor,
offset: Union[int, torch.Tensor] = 0
self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Input x.
@@ -894,16 +891,14 @@ def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
"""
batch_size = lengths.size(0)
max_len = max_len if max_len > 0 else lengths.max().item()
seq_range = torch.arange(0,
max_len,
dtype=torch.int64,
device=lengths.device)
seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device)
seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
seq_length_expand = lengths.unsqueeze(-1)
mask = seq_range_expand >= seq_length_expand
return mask
#https://github.com/FunAudioLLM/CosyVoice/blob/main/examples/magicdata-read/cosyvoice/conf/cosyvoice.yaml
# https://github.com/FunAudioLLM/CosyVoice/blob/main/examples/magicdata-read/cosyvoice/conf/cosyvoice.yaml
class ConformerEncoder(torch.nn.Module):
"""Conformer encoder module."""
@@ -917,14 +912,14 @@ class ConformerEncoder(torch.nn.Module):
dropout_rate: float = 0.1,
positional_dropout_rate: float = 0.1,
attention_dropout_rate: float = 0.0,
input_layer: str = 'linear',
pos_enc_layer_type: str = 'rel_pos_espnet',
input_layer: str = "linear",
pos_enc_layer_type: str = "rel_pos_espnet",
normalize_before: bool = True,
static_chunk_size: int = 1, # 1: causal_mask; 0: full_mask
static_chunk_size: int = 1, # 1: causal_mask; 0: full_mask
use_dynamic_chunk: bool = False,
use_dynamic_left_chunk: bool = False,
positionwise_conv_kernel_size: int = 1,
macaron_style: bool =False,
macaron_style: bool = False,
selfattention_layer_type: str = "rel_selfattn",
activation_type: str = "swish",
use_cnn_module: bool = False,
@@ -953,13 +948,17 @@ class ConformerEncoder(torch.nn.Module):
"""
super().__init__()
self.output_size = output_size
self.embed = LinearEmbed(input_size, output_size, dropout_rate,
EspnetRelPositionalEncoding(output_size, positional_dropout_rate))
self.embed = LinearEmbed(
input_size,
output_size,
dropout_rate,
EspnetRelPositionalEncoding(output_size, positional_dropout_rate),
)
self.normalize_before = normalize_before
self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5)
self.gradient_checkpointing = gradient_checkpointing
self.use_dynamic_chunk = use_dynamic_chunk
self.static_chunk_size = static_chunk_size
self.use_dynamic_chunk = use_dynamic_chunk
self.use_dynamic_left_chunk = use_dynamic_left_chunk
@@ -980,40 +979,60 @@ class ConformerEncoder(torch.nn.Module):
activation,
)
# convolution module definition
convolution_layer_args = (output_size, cnn_module_kernel, activation,
cnn_module_norm, causal)
convolution_layer_args = (
output_size,
cnn_module_kernel,
activation,
cnn_module_norm,
causal,
)
self.encoders = torch.nn.ModuleList([
ConformerEncoderLayer(
output_size,
RelPositionMultiHeadedAttention(
*encoder_selfattn_layer_args),
PositionwiseFeedForward(*positionwise_layer_args),
PositionwiseFeedForward(
*positionwise_layer_args) if macaron_style else None,
ConvolutionModule(
*convolution_layer_args) if use_cnn_module else None,
dropout_rate,
normalize_before,
) for _ in range(num_blocks)
])
def forward_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor,
self.encoders = torch.nn.ModuleList(
[
ConformerEncoderLayer(
output_size,
RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args),
PositionwiseFeedForward(*positionwise_layer_args),
(
PositionwiseFeedForward(*positionwise_layer_args)
if macaron_style
else None
),
(
ConvolutionModule(*convolution_layer_args)
if use_cnn_module
else None
),
dropout_rate,
normalize_before,
)
for _ in range(num_blocks)
]
)
def forward_layers(
self,
xs: torch.Tensor,
chunk_masks: torch.Tensor,
pos_emb: torch.Tensor,
mask_pad: torch.Tensor) -> torch.Tensor:
mask_pad: torch.Tensor,
) -> torch.Tensor:
for layer in self.encoders:
xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad)
return xs
@torch.jit.unused
def forward_layers_checkpointed(self, xs: torch.Tensor,
chunk_masks: torch.Tensor,
pos_emb: torch.Tensor,
mask_pad: torch.Tensor) -> torch.Tensor:
def forward_layers_checkpointed(
self,
xs: torch.Tensor,
chunk_masks: torch.Tensor,
pos_emb: torch.Tensor,
mask_pad: torch.Tensor,
) -> torch.Tensor:
for layer in self.encoders:
xs, chunk_masks, _, _ = ckpt.checkpoint(layer.__call__, xs,
chunk_masks, pos_emb,
mask_pad)
xs, chunk_masks, _, _ = ckpt.checkpoint(
layer.__call__, xs, chunk_masks, pos_emb, mask_pad
)
return xs
def forward(
@@ -1047,18 +1066,20 @@ class ConformerEncoder(torch.nn.Module):
https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690/2
"""
T = xs.size(1)
masks = pad_mask.to(torch.bool).unsqueeze(1) # (B, 1, T)
masks = pad_mask.to(torch.bool).unsqueeze(1) # (B, 1, T)
xs, pos_emb = self.embed(xs)
mask_pad = masks # (B, 1, T/subsample_rate)
chunk_masks = add_optional_chunk_mask(xs, masks,
self.use_dynamic_chunk,
self.use_dynamic_left_chunk,
decoding_chunk_size,
self.static_chunk_size,
num_decoding_left_chunks)
chunk_masks = add_optional_chunk_mask(
xs,
masks,
self.use_dynamic_chunk,
self.use_dynamic_left_chunk,
decoding_chunk_size,
self.static_chunk_size,
num_decoding_left_chunks,
)
if self.gradient_checkpointing and self.training:
xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb,
mask_pad)
xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb, mask_pad)
else:
xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad)
if self.normalize_before:
@@ -1067,4 +1088,3 @@ class ConformerEncoder(torch.nn.Module):
# return the masks before encoder layers, and the masks will be used
# for cross attention with decoder later
return xs, masks
acestep/models/lyrics_utils/lyric_normalizer.py
+28 -25
View File
@@ -2,39 +2,42 @@ import re
from opencc import OpenCC
t2s_converter = OpenCC('t2s')
s2t_converter = OpenCC('s2t')
t2s_converter = OpenCC("t2s")
s2t_converter = OpenCC("s2t")
EMOJI_PATTERN = re.compile(
"["
"\U0001F600-\U0001F64F" # Emoticons
"]+", flags=re.UNICODE
"\U0001f600-\U0001f64f" # Emoticons
"]+",
flags=re.UNICODE,
)
# 创建一个翻译表，用于替换和移除字符
TRANSLATION_TABLE = str.maketrans({
'-': ' ', # 将 '-' 替换为空格
',': None,
'.': None,
'，': None,
'。': None,
'!': None,
'！': None,
'?': None,
'？': None,
'…': None,
';': None,
'；': None,
':': None,
'：': None,
'\u3000': ' ', # 将全角空格替换为空格
})
TRANSLATION_TABLE = str.maketrans(
{
"-": " ", # 将 '-' 替换为空格
",": None,
".": None,
"，": None,
"。": None,
"!": None,
"！": None,
"?": None,
"？": None,
"…": None,
";": None,
"；": None,
":": None,
"：": None,
"\u3000": " ", # 将全角空格替换为空格
}
)
# 替换括号中的内容，包括中括号和小括号
BACKSLASH_PATTERN = re.compile(r'\(.*?\)|\[.*?\]')
BACKSLASH_PATTERN = re.compile(r"\(.*?\)|\[.*?\]")
SPACE_PATTERN = re.compile('(?<!^)\s+(?!$)')
SPACE_PATTERN = re.compile("(?<!^)\s+(?!$)")
def normalize_text(text, language, strip=True):
@@ -45,10 +48,10 @@ def normalize_text(text, language, strip=True):
text = text.translate(TRANSLATION_TABLE)
# Step 2: 移除表情符号
text = EMOJI_PATTERN.sub('', text)
text = EMOJI_PATTERN.sub("", text)
# Step 3: 连续空白字符替换为单个空格，首位除外
text = SPACE_PATTERN.sub(' ', text)
text = SPACE_PATTERN.sub(" ", text)
# Step 4: 去除首尾空白字符（如果需要）
if strip:
acestep/models/lyrics_utils/lyric_tokenizer.py
+143 -39
View File
@@ -19,7 +19,7 @@ from .zh_num2words import TextNorm as zh_num2words
from typing import Dict, List, Optional, Set, Union
#copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/tokenizer.py
# copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/tokenizer.py
def get_spacy_lang(lang):
if lang == "zh":
return Chinese()
@@ -446,7 +446,9 @@ _ordinal_re = {
"it": re.compile(r"([0-9]+)(º|°|ª|o|a|i|e)"),
"pl": re.compile(r"([0-9]+)(º|ª|st|nd|rd|th)"),
"ar": re.compile(r"([0-9]+)(ون|ين|ث|ر|ى)"),
"cs": re.compile(r"([0-9]+)\.(?=\s|$)"), # In Czech, a dot is often used after the number to indicate ordinals.
"cs": re.compile(
r"([0-9]+)\.(?=\s|$)"
), # In Czech, a dot is often used after the number to indicate ordinals.
"ru": re.compile(r"([0-9]+)(-й|-я|-е|-ое|-ье|-го)"),
"nl": re.compile(r"([0-9]+)(de|ste|e)"),
"tr": re.compile(r"([0-9]+)(\.|inci|nci|uncu|üncü|\.)"),
@@ -486,7 +488,9 @@ def _expand_decimal_point(m, lang="en"):
def _expand_currency(m, lang="en", currency="USD"):
amount = float((re.sub(r"[^\d.]", "", m.group(0).replace(",", "."))))
full_amount = num2words(amount, to="currency", currency=currency, lang=lang if lang != "cs" else "cz")
full_amount = num2words(
amount, to="currency", currency=currency, lang=lang if lang != "cs" else "cz"
)
and_equivalents = {
"en": ", ",
@@ -530,13 +534,21 @@ def expand_numbers_multilingual(text, lang="en"):
else:
text = re.sub(_dot_number_re, _remove_dots, text)
try:
text = re.sub(_currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text)
text = re.sub(_currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text)
text = re.sub(_currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text)
text = re.sub(
_currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text
)
text = re.sub(
_currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text
)
text = re.sub(
_currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text
)
except:
pass
if lang != "tr":
text = re.sub(_decimal_number_re, lambda m: _expand_decimal_point(m, lang), text)
text = re.sub(
_decimal_number_re, lambda m: _expand_decimal_point(m, lang), text
)
text = re.sub(_ordinal_re[lang], lambda m: _expand_ordinal(m, lang), text)
text = re.sub(_number_re, lambda m: _expand_number(m, lang), text)
return text
@@ -582,7 +594,15 @@ def basic_cleaners(text):
def chinese_transliterate(text):
return "".join(
[p[0] for p in pypinyin.pinyin(text, style=pypinyin.Style.TONE3, heteronym=False, neutral_tone_with_five=True)]
[
p[0]
for p in pypinyin.pinyin(
text,
style=pypinyin.Style.TONE3,
heteronym=False,
neutral_tone_with_five=True,
)
]
)
@@ -597,7 +617,9 @@ def korean_transliterate(text):
return r.translit(text)
DEFAULT_VOCAB_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "vocab.json")
DEFAULT_VOCAB_FILE = os.path.join(
os.path.dirname(os.path.realpath(__file__)), "vocab.json"
)
class VoiceBpeTokenizer:
@@ -639,7 +661,23 @@ class VoiceBpeTokenizer:
# )
def preprocess_text(self, txt, lang):
if lang in {"ar", "cs", "de", "en", "es", "fr", "hu", "it", "nl", "pl", "pt", "ru", "tr", "zh", "ko"}:
if lang in {
"ar",
"cs",
"de",
"en",
"es",
"fr",
"hu",
"it",
"nl",
"pl",
"pt",
"ru",
"tr",
"zh",
"ko",
}:
txt = multilingual_cleaners(txt, lang)
if lang == "zh":
txt = chinese_transliterate(txt)
@@ -653,7 +691,7 @@ class VoiceBpeTokenizer:
else:
raise NotImplementedError(f"Language '{lang}' is not supported.")
return txt
def encode(self, txt, lang):
lang = lang.split("-")[0] # remove the region
self.check_input_length(txt, lang)
@@ -671,12 +709,13 @@ class VoiceBpeTokenizer:
txt = txt.replace("[STOP]", "")
# txt = txt.replace("[UNK]", "")
return txt
#copy from https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3936
# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3936
def batch_decode(
self,
sequences: Union[List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"],
sequences: Union[
List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"
],
skip_special_tokens: bool = False,
) -> List[str]:
"""
@@ -693,13 +732,10 @@ class VoiceBpeTokenizer:
Returns:
`List[str]`: The list of decoded sentences.
"""
return [
self.decode(seq)
for seq in sequences
]
#https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/layers/xtts/trainer/dataset.py#L202
# def pad(self):
return [self.decode(seq) for seq in sequences]
# https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/layers/xtts/trainer/dataset.py#L202
# def pad(self):
def __len__(self):
return self.tokenizer.get_vocab_size()
@@ -716,15 +752,27 @@ def test_expand_numbers_multilingual():
("This is a 1st test", "This is a first test", "en"),
("That will be $20 sir.", "That will be twenty dollars sir.", "en"),
("That will be 20€ sir.", "That will be twenty euro sir.", "en"),
("That will be 20.15€ sir.", "That will be twenty euro, fifteen cents sir.", "en"),
(
"That will be 20.15€ sir.",
"That will be twenty euro, fifteen cents sir.",
"en",
),
("That's 100,000.5.", "That's one hundred thousand point five.", "en"),
# French
("En 12,5 secondes.", "En douze virgule cinq secondes.", "fr"),
("Il y avait 50 soldats.", "Il y avait cinquante soldats.", "fr"),
("Ceci est un 1er test", "Ceci est un premier test", "fr"),
("Cela vous fera $20 monsieur.", "Cela vous fera vingt dollars monsieur.", "fr"),
(
"Cela vous fera $20 monsieur.",
"Cela vous fera vingt dollars monsieur.",
"fr",
),
("Cela vous fera 20€ monsieur.", "Cela vous fera vingt euros monsieur.", "fr"),
("Cela vous fera 20,15€ monsieur.", "Cela vous fera vingt euros et quinze centimes monsieur.", "fr"),
(
"Cela vous fera 20,15€ monsieur.",
"Cela vous fera vingt euros et quinze centimes monsieur.",
"fr",
),
("Ce sera 100.000,5.", "Ce sera cent mille virgule cinq.", "fr"),
# German
("In 12,5 Sekunden.", "In zwölf Komma fünf Sekunden.", "de"),
@@ -732,21 +780,33 @@ def test_expand_numbers_multilingual():
("Dies ist ein 1. Test", "Dies ist ein erste Test", "de"), # Issue with gender
("Das macht $20 Herr.", "Das macht zwanzig Dollar Herr.", "de"),
("Das macht 20€ Herr.", "Das macht zwanzig Euro Herr.", "de"),
("Das macht 20,15€ Herr.", "Das macht zwanzig Euro und fünfzehn Cent Herr.", "de"),
(
"Das macht 20,15€ Herr.",
"Das macht zwanzig Euro und fünfzehn Cent Herr.",
"de",
),
# Spanish
("En 12,5 segundos.", "En doce punto cinco segundos.", "es"),
("Había 50 soldados.", "Había cincuenta soldados.", "es"),
("Este es un 1er test", "Este es un primero test", "es"),
("Eso le costará $20 señor.", "Eso le costará veinte dólares señor.", "es"),
("Eso le costará 20€ señor.", "Eso le costará veinte euros señor.", "es"),
("Eso le costará 20,15€ señor.", "Eso le costará veinte euros con quince céntimos señor.", "es"),
(
"Eso le costará 20,15€ señor.",
"Eso le costará veinte euros con quince céntimos señor.",
"es",
),
# Italian
("In 12,5 secondi.", "In dodici virgola cinque secondi.", "it"),
("C'erano 50 soldati.", "C'erano cinquanta soldati.", "it"),
("Questo è un 1° test", "Questo è un primo test", "it"),
("Ti costerà $20 signore.", "Ti costerà venti dollari signore.", "it"),
("Ti costerà 20€ signore.", "Ti costerà venti euro signore.", "it"),
("Ti costerà 20,15€ signore.", "Ti costerà venti euro e quindici centesimi signore.", "it"),
(
"Ti costerà 20,15€ signore.",
"Ti costerà venti euro e quindici centesimi signore.",
"it",
),
# Portuguese
("Em 12,5 segundos.", "Em doze vírgula cinco segundos.", "pt"),
("Havia 50 soldados.", "Havia cinquenta soldados.", "pt"),
@@ -761,8 +821,16 @@ def test_expand_numbers_multilingual():
# Polish
("W 12,5 sekundy.", "W dwanaście przecinek pięć sekundy.", "pl"),
("Było 50 żołnierzy.", "Było pięćdziesiąt żołnierzy.", "pl"),
("To będzie kosztować 20€ panie.", "To będzie kosztować dwadzieścia euro panie.", "pl"),
("To będzie kosztować 20,15€ panie.", "To będzie kosztować dwadzieścia euro, piętnaście centów panie.", "pl"),
(
"To będzie kosztować 20€ panie.",
"To będzie kosztować dwadzieścia euro panie.",
"pl",
),
(
"To będzie kosztować 20,15€ panie.",
"To będzie kosztować dwadzieścia euro, piętnaście centów panie.",
"pl",
),
# Arabic
("في الـ 12,5 ثانية.", "في الـ اثنا عشر , خمسون ثانية.", "ar"),
("كان هناك 50 جنديًا.", "كان هناك خمسون جنديًا.", "ar"),
@@ -776,8 +844,16 @@ def test_expand_numbers_multilingual():
# Russian
("Через 12.5 секунды.", "Через двенадцать запятая пять секунды.", "ru"),
("Там было 50 солдат.", "Там было пятьдесят солдат.", "ru"),
("Это будет 20.15€ сэр.", "Это будет двадцать евро, пятнадцать центов сэр.", "ru"),
("Это будет стоить 20€ господин.", "Это будет стоить двадцать евро господин.", "ru"),
(
"Это будет 20.15€ сэр.",
"Это будет двадцать евро, пятнадцать центов сэр.",
"ru",
),
(
"Это будет стоить 20€ господин.",
"Это будет стоить двадцать евро господин.",
"ru",
),
# Dutch
("In 12,5 seconden.", "In twaalf komma vijf seconden.", "nl"),
("Er waren 50 soldaten.", "Er waren vijftig soldaten.", "nl"),
@@ -817,18 +893,30 @@ def test_abbreviations_multilingual():
("La Dra. Martinez es muy buena.", "La doctora Martinez es muy buena.", "es"),
# French
("Bonjour Mr. Dupond.", "Bonjour monsieur Dupond.", "fr"),
("Mme. Moreau est absente aujourd'hui.", "madame Moreau est absente aujourd'hui.", "fr"),
(
"Mme. Moreau est absente aujourd'hui.",
"madame Moreau est absente aujourd'hui.",
"fr",
),
# German
("Frau Dr. Müller ist sehr klug.", "Frau doktor Müller ist sehr klug.", "de"),
# Portuguese
("Olá Sr. Silva.", "Olá senhor Silva.", "pt"),
("Dra. Costa, você está disponível?", "doutora Costa, você está disponível?", "pt"),
(
"Dra. Costa, você está disponível?",
"doutora Costa, você está disponível?",
"pt",
),
# Italian
("Buongiorno, Sig. Rossi.", "Buongiorno, signore Rossi.", "it"),
# ("Sig.ra Bianchi, posso aiutarti?", 'signora Bianchi, posso aiutarti?', 'it'), # Issue with matching that pattern
# Polish
("Dzień dobry, P. Kowalski.", "Dzień dobry, pani Kowalski.", "pl"),
("M. Nowak, czy mogę zadać pytanie?", "pan Nowak, czy mogę zadać pytanie?", "pl"),
(
"M. Nowak, czy mogę zadać pytanie?",
"pan Nowak, czy mogę zadać pytanie?",
"pl",
),
# Czech
("P. Novák", "pan Novák", "cs"),
("Dr. Vojtěch", "doktor Vojtěch", "cs"),
@@ -837,7 +925,11 @@ def test_abbreviations_multilingual():
("Mevr. de Vries", "mevrouw de Vries", "nl"),
# Russian
("Здравствуйте Г-н Иванов.", "Здравствуйте господин Иванов.", "ru"),
("Д-р Смирнов здесь, чтобы увидеть вас.", "доктор Смирнов здесь, чтобы увидеть вас.", "ru"),
(
"Д-р Смирнов здесь, чтобы увидеть вас.",
"доктор Смирнов здесь, чтобы увидеть вас.",
"ru",
),
# Turkish
("Merhaba B. Yılmaz.", "Merhaba bay Yılmaz.", "tr"),
("Dr. Ayşe burada.", "doktor Ayşe burada.", "tr"),
@@ -856,8 +948,16 @@ def test_symbols_multilingual():
("Te veo @ la fiesta", "Te veo arroba la fiesta", "es"),
("J'ai 14° de fièvre", "J'ai 14 degrés de fièvre", "fr"),
("Die Rechnung beträgt £ 20", "Die Rechnung beträgt pfund 20", "de"),
("O meu email é ana&joao@gmail.com", "O meu email é ana e joao arroba gmail.com", "pt"),
("linguaggio di programmazione C#", "linguaggio di programmazione C cancelletto", "it"),
(
"O meu email é ana&joao@gmail.com",
"O meu email é ana e joao arroba gmail.com",
"pt",
),
(
"linguaggio di programmazione C#",
"linguaggio di programmazione C cancelletto",
"it",
),
("Moja temperatura to 36.6°", "Moja temperatura to 36.6 stopnie", "pl"),
("Mám 14% baterie", "Mám 14 procento baterie", "cs"),
("Těším se na tebe @ party", "Těším se na tebe na party", "cs"),
@@ -868,7 +968,11 @@ def test_symbols_multilingual():
("لدي 14% في البطارية", "لدي 14 في المئة في البطارية", "ar"),
("我的电量为 14%", "我的电量为 14 百分之", "zh"),
("Pilim %14 dolu.", "Pilim yüzde 14 dolu.", "tr"),
("Az akkumulátorom töltöttsége 14%", "Az akkumulátorom töltöttsége 14 százalék", "hu"),
(
"Az akkumulátorom töltöttsége 14%",
"Az akkumulátorom töltöttsége 14 százalék",
"hu",
),
("배터리 잔량이 14%입니다.", "배터리 잔량이 14 퍼센트입니다.", "ko"),
]
@@ -880,4 +984,4 @@ def test_symbols_multilingual():
if __name__ == "__main__":
test_expand_numbers_multilingual()
test_abbreviations_multilingual()
test_symbols_multilingual()
test_symbols_multilingual()
acestep/models/lyrics_utils/zh_num2words.py
+2 -2
View File
@@ -1,7 +1,7 @@
# Authors:
# 2019.5 Zhiyang Zhou (https://github.com/Joee1995/chn_text_norm.git)
# 2019.9 - 2022 Jiayu DU
#copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/zh_num2words.py
# copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/zh_num2words.py
import argparse
import csv
import os
@@ -1206,4 +1206,4 @@ if __name__ == "__main__":
ndone += 1
if ndone % args.log_interval == 0:
print(f"text norm: {ndone} lines done.", file=sys.stderr, flush=True)
print(f"text norm: {ndone} lines done in total.", file=sys.stderr, flush=True)
print(f"text norm: {ndone} lines done in total.", file=sys.stderr, flush=True)