all inference code
This commit is contained in:
sean
@@ -0,0 +1,412 @@
|
||||
import torch.nn as nn
|
||||
import torch
|
||||
import random
|
||||
from loguru import logger
|
||||
|
||||
from transformers import UMT5EncoderModel, AutoTokenizer, AutoModel
|
||||
import re
|
||||
from typing import List, Tuple, Dict, Set
|
||||
import sys
|
||||
import os
|
||||
|
||||
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
|
||||
class TrieNode:
|
||||
def __init__(self):
|
||||
self.children: Dict[str, 'TrieNode'] = {}
|
||||
self.is_end_of_word: bool = False
|
||||
|
||||
|
||||
class Trie:
|
||||
def __init__(self):
|
||||
self.root = TrieNode()
|
||||
|
||||
def insert(self, word: str):
|
||||
node = self.root
|
||||
for char in word:
|
||||
if char not in node.children:
|
||||
node.children[char] = TrieNode()
|
||||
node = node.children[char]
|
||||
node.is_end_of_word = True
|
||||
|
||||
def search_from(self, s: str, start: int) -> List[str]:
|
||||
"""
|
||||
从字符串 s 的位置 start 开始，使用 Trie 树查找所有可能的匹配 phoneme。
|
||||
返回所有匹配的 phoneme。
|
||||
"""
|
||||
node = self.root
|
||||
matches = []
|
||||
current_phoneme = []
|
||||
for i in range(start, len(s)):
|
||||
char = s[i]
|
||||
if char in node.children:
|
||||
node = node.children[char]
|
||||
current_phoneme.append(char)
|
||||
if node.is_end_of_word:
|
||||
matches.append(''.join(current_phoneme))
|
||||
else:
|
||||
break
|
||||
return matches
|
||||
|
||||
|
||||
class PhonemeMatcher:
|
||||
def __init__(self, word_dict: Set[str]):
|
||||
"""
|
||||
初始化 PhonemeMatcher，构建 Trie 树。
|
||||
|
||||
:param word_dict: Set[str] - 包含所有 phoneme 的集合
|
||||
"""
|
||||
self.trie = Trie()
|
||||
for word in word_dict:
|
||||
self.trie.insert(word)
|
||||
|
||||
def tokenize(self, s: str) -> List[str]:
|
||||
"""
|
||||
将输入的 xsampa 字符串拆分成 phoneme 序列，尽可能使用词表中的 phoneme，
|
||||
并在无法完全匹配时，选择编辑距离最小且 phoneme 数量最少的序列。
|
||||
|
||||
:param s: str - 输入的 xsampa 字符串
|
||||
:return: List[str] - 输出的 phoneme 序列
|
||||
"""
|
||||
n = len(s)
|
||||
# 初始化 DP 数组，dp[i] = (cost, phoneme_count, phone_list)
|
||||
dp: List[Tuple[int, int, List[str]]] = [(sys.maxsize, sys.maxsize, []) for _ in range(n + 1)]
|
||||
dp[0] = (0, 0, [])
|
||||
|
||||
for i in range(n):
|
||||
current_cost, current_count, current_list = dp[i]
|
||||
if current_cost == sys.maxsize:
|
||||
continue # 无法到达当前位置
|
||||
|
||||
# 查找所有从位置 i 开始的匹配 phoneme
|
||||
matches = self.trie.search_from(s, i)
|
||||
|
||||
if matches:
|
||||
for phoneme in matches:
|
||||
end = i + len(phoneme)
|
||||
new_cost = current_cost # 匹配成功，无需增加编辑距离
|
||||
new_count = current_count + 1
|
||||
new_list = current_list + [phoneme]
|
||||
|
||||
if new_cost < dp[end][0]:
|
||||
dp[end] = (new_cost, new_count, new_list)
|
||||
elif new_cost == dp[end][0]:
|
||||
if new_count < dp[end][1]:
|
||||
dp[end] = (new_cost, new_count, new_list)
|
||||
else:
|
||||
# 没有匹配的 phoneme，考虑跳过当前字符，增加编辑距离
|
||||
new_cost = current_cost + 1
|
||||
end = i + 1
|
||||
new_count = current_count + 1 # 跳过一个字符也算作一个 phoneme
|
||||
new_list = current_list + [s[i]]
|
||||
|
||||
if new_cost < dp[end][0]:
|
||||
dp[end] = (new_cost, new_count, new_list)
|
||||
elif new_cost == dp[end][0]:
|
||||
if new_count < dp[end][1]:
|
||||
dp[end] = (new_cost, new_count, new_list)
|
||||
|
||||
# 如果无法完全匹配，选择最优的近似匹配
|
||||
if dp[n][0] == sys.maxsize:
|
||||
# 找到所有可能的最小编辑距离
|
||||
min_cost = min(dp[i][0] for i in range(n + 1))
|
||||
# 选择最小编辑距离且 phoneme 数量最少的序列
|
||||
candidates = [dp[i] for i in range(n + 1) if dp[i][0] == min_cost]
|
||||
if candidates:
|
||||
# 选择 phoneme 数量最少的
|
||||
best = min(candidates, key=lambda x: x[1])
|
||||
return best[2]
|
||||
else:
|
||||
return []
|
||||
|
||||
return dp[n][2]
|
||||
|
||||
|
||||
HARMONIX_LABELS = [
|
||||
'start',
|
||||
'end',
|
||||
'intro',
|
||||
'outro',
|
||||
'break',
|
||||
'bridge',
|
||||
'inst',
|
||||
'solo',
|
||||
'verse',
|
||||
'chorus',
|
||||
]
|
||||
|
||||
|
||||
def timestamp2second(timestamps):
|
||||
res = []
|
||||
for item in timestamps:
|
||||
start, end = item["start"], item["end"]
|
||||
# convert 8kHz to latents level
|
||||
start = round(start / 8000, 2)
|
||||
end = round(end / 8000, 2)
|
||||
res.append({"start": start, "end": end})
|
||||
return res
|
||||
|
||||
|
||||
def sample_lyric_mask(voiced_timestamp, max_length):
|
||||
voiced_timestamps = timestamp2second(voiced_timestamp)
|
||||
|
||||
min_gaps = [1,2,3,4,5]
|
||||
while len(min_gaps) > 0:
|
||||
min_gap = min_gaps.pop()
|
||||
can_split_breaks = []
|
||||
last_end = 0.00
|
||||
for item in voiced_timestamps:
|
||||
if item["start"] - last_end >= min_gap:
|
||||
if last_end == 0.00:
|
||||
can_split_breaks.append((last_end, item["start"] - 0.5))
|
||||
else:
|
||||
can_split_breaks.append((last_end + 0.5, item["start"] - 0.5))
|
||||
last_end = item["end"]
|
||||
if len(can_split_breaks) > 1:
|
||||
if can_split_breaks[1][0] <= 360:
|
||||
break
|
||||
else:
|
||||
if min_gap == 1:
|
||||
return 0.0, 360.0, 36
|
||||
|
||||
if len(can_split_breaks) == 0:
|
||||
mask_start, mask_end = 0.0, max_length
|
||||
min_cut_level = int(mask_end//10 - mask_start//10 + 1)
|
||||
return 0.0, mask_end, min_cut_level
|
||||
|
||||
if len(can_split_breaks) == 1:
|
||||
# 前后随机选一个
|
||||
mask_start = random.choice(["start", "middle"])
|
||||
if mask_start == "start":
|
||||
mask_start = 0.0
|
||||
mask_end = random.uniform(can_split_breaks[0][0], can_split_breaks[0][1])
|
||||
else:
|
||||
mask_start = random.uniform(can_split_breaks[0][0], can_split_breaks[0][1])
|
||||
mask_end = max_length
|
||||
min_cut_level = int(mask_end//10 - mask_start//10 + 1)
|
||||
return mask_start, mask_end, min_cut_level
|
||||
|
||||
mask_start, mask_end = 0.0, 370
|
||||
min_cut_level = 37
|
||||
breaths_gap = [end-start for start, end in can_split_breaks]
|
||||
max_tried = 5
|
||||
while mask_end - mask_start > 370 and min_cut_level > 0 and min_cut_level > 36:
|
||||
total_breaths = len(can_split_breaks)
|
||||
start = random.choices(range(total_breaths-1), weights=breaths_gap[:-1])[0]
|
||||
end = random.choices(range(start + 1, total_breaths), weights=breaths_gap[start+1:], k=1)[0]
|
||||
start_break, end_break = can_split_breaks[start], can_split_breaks[end]
|
||||
mask_start, mask_end = random.uniform(start_break[0], start_break[1]), random.uniform(end_break[0], end_break[1])
|
||||
min_cut_level = int(mask_end//10 - mask_start//10 + 1)
|
||||
if min_cut_level < 36:
|
||||
min_cut_level = random.randint(min_cut_level, 36)
|
||||
if max_tried == 0:
|
||||
print("max tried", mask_start, mask_end, min_cut_level, "breaths_gap", breaths_gap, "can_split_breaks", can_split_breaks)
|
||||
break
|
||||
max_tried -= 1
|
||||
mask_start, mask_end = round(mask_start, 2), min(round(mask_end, 2), max_length)
|
||||
return mask_start, mask_end, min_cut_level
|
||||
|
||||
|
||||
def check_valid_lyric_lines(lyric_lines):
|
||||
# must has lyric lines
|
||||
if len(lyric_lines) == 0:
|
||||
return False
|
||||
for valid_lyric_line in lyric_lines:
|
||||
if len(valid_lyric_line[1]) > 0:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def select_valid_lyric_lines(lyric_lines, mask_start, mask_end):
|
||||
# 选歌词原则
|
||||
# 宁可多，不可少
|
||||
# 选取mask_start和mask_end之间的歌词行，如果mask_end在一个歌词行中间，那么这个歌词行也要被选取，但最后的structure不要
|
||||
valid_lyric_lines = []
|
||||
add_tail_structure = True
|
||||
for lyric_line in lyric_lines:
|
||||
if lyric_line["start"] > lyric_line["end"]:
|
||||
continue
|
||||
if lyric_line["start"]+1.0 >= mask_start and lyric_line["end"]-1.0 <= mask_end:
|
||||
if len(valid_lyric_lines) > 0:
|
||||
if valid_lyric_lines[-1][0] is not None and valid_lyric_lines[-1][0] != lyric_line["structure"] and lyric_line["structure"] != "":
|
||||
valid_lyric_lines.append((lyric_line["structure"], [], [], (lyric_line["start"], lyric_line["end"])))
|
||||
elif lyric_line["structure"] != "":
|
||||
valid_lyric_lines.append((lyric_line["structure"], [], [], (lyric_line["start"], lyric_line["end"])))
|
||||
lyric_line["lyric_line"] = lyric_line["lyric_line"].strip()
|
||||
if lyric_line["lyric_line"] and "phoneme_line_ipa" in lyric_line and len(lyric_line["phoneme_line_ipa"]) > 0:
|
||||
valid_lyric_lines.append((None, lyric_line["lyric_line"], lyric_line["phoneme_line_ipa"], (lyric_line["start"], lyric_line["end"])))
|
||||
elif mask_start < lyric_line["start"] and lyric_line["start"] < mask_end and lyric_line["end"] > mask_end:
|
||||
lyric_line["lyric_line"] = lyric_line["lyric_line"].strip()
|
||||
if lyric_line["lyric_line"] and "phoneme_line_ipa" in lyric_line and len(lyric_line["phoneme_line_ipa"]) > 0:
|
||||
valid_lyric_lines.append((None, lyric_line["lyric_line"], lyric_line["phoneme_line_ipa"], (lyric_line["start"], lyric_line["end"])))
|
||||
add_tail_structure = False
|
||||
break
|
||||
elif lyric_line["start"] > mask_start and lyric_line["start"] < mask_end and not lyric_line["lyric_line"] and add_tail_structure:
|
||||
valid_lyric_lines.append((lyric_line["structure"], [], [], (lyric_line["start"], lyric_line["end"])))
|
||||
add_tail_structure = False
|
||||
break
|
||||
if len(valid_lyric_lines) > 0 and len(lyric_lines) > 0 and add_tail_structure:
|
||||
if lyric_lines[-1]["structure"] != "" and lyric_lines[-1]["structure"] != valid_lyric_lines[-1][0]:
|
||||
if lyric_lines[-1]["start"] > mask_start and lyric_lines[-1]["start"] < mask_end:
|
||||
valid_lyric_lines.append((lyric_lines[-1]["structure"], [], [], (lyric_lines[-1]["start"], lyric_lines[-1]["end"])))
|
||||
return valid_lyric_lines
|
||||
|
||||
|
||||
def sample_lyric_mask_with_cut_levels(voiced_timestamp, cut_level, n_chunks, lyric_lines):
|
||||
voiced_timestamps = timestamp2second(voiced_timestamp)
|
||||
|
||||
candidate_spans = []
|
||||
for candidate_start_idx in range(n_chunks):
|
||||
candidate_start_second = candidate_start_idx * 10
|
||||
candidate_end_second = (candidate_start_idx + cut_level) * 10
|
||||
valid = True
|
||||
for item in voiced_timestamps:
|
||||
if item["start"] < candidate_start_second and candidate_start_second < item["end"]:
|
||||
valid = False
|
||||
break
|
||||
if item["start"] < candidate_end_second and candidate_end_second < item["end"]:
|
||||
valid = False
|
||||
break
|
||||
valid_lyric_lines = select_valid_lyric_lines(lyric_lines, candidate_start_second, candidate_end_second)
|
||||
if not check_valid_lyric_lines(valid_lyric_lines):
|
||||
valid = False
|
||||
if valid:
|
||||
candidate_spans.append((candidate_start_second, candidate_end_second, valid_lyric_lines))
|
||||
|
||||
if len(candidate_spans) > 0:
|
||||
return candidate_spans
|
||||
else:
|
||||
candidate_spans = []
|
||||
for candidate_start_idx in range(n_chunks):
|
||||
candidate_start_second = candidate_start_idx * 10
|
||||
candidate_end_second = (candidate_start_idx + cut_level) * 10
|
||||
valid_lyric_lines = select_valid_lyric_lines(lyric_lines, candidate_start_second, candidate_end_second)
|
||||
if check_valid_lyric_lines(valid_lyric_lines):
|
||||
candidate_spans.append((candidate_start_second, candidate_end_second, valid_lyric_lines))
|
||||
if len(candidate_spans) > 0:
|
||||
return candidate_spans
|
||||
return []
|
||||
|
||||
|
||||
def sample_lyric_mask_with_lyric_timestamp(cut_level, lyric_lines, expected_num_example, n_chunks, start_pad_offset=1.0):
|
||||
# 1 去掉structure
|
||||
# non_structure_lyric_lines = [lyric_line for lyric_line in lyric_lines if lyric_line["lyric_line"] and "phoneme_line_ipa" in lyric_line and len(lyric_line["phoneme_line_ipa"]) > 0 and lyric_line["start"] < lyric_line["end"]]
|
||||
# 保留structure
|
||||
valid_lyric_lines = []
|
||||
last_structure = ""
|
||||
for lyric_line in lyric_lines:
|
||||
if "structure" not in lyric_line:
|
||||
lyric_line["structure"] = ""
|
||||
if lyric_line["start"] < lyric_line["end"]:
|
||||
new_line = lyric_line.copy()
|
||||
if not lyric_line["lyric_line"] or "phoneme_line_ipa" not in lyric_line or len(lyric_line["phoneme_line_ipa"]) == 0:
|
||||
if lyric_line["structure"] != "":
|
||||
new_line["lyric_line"] = "["+lyric_line["structure"]+"]"
|
||||
new_line["phoneme_line_ipa"] = ["_"]
|
||||
else:
|
||||
last_structure = lyric_line["structure"]
|
||||
continue
|
||||
else:
|
||||
if new_line["structure"] != "" and new_line["structure"] != last_structure:
|
||||
if new_line["lyric_line"] != "[" + new_line["structure"] + "]":
|
||||
new_line["lyric_line"] = f"[{new_line['structure']}]\n{new_line['lyric_line']}"
|
||||
new_line["phoneme_line_ipa"] = ["_", "_"] + new_line["phoneme_line_ipa"]
|
||||
|
||||
valid_lyric_lines.append(new_line)
|
||||
last_structure = lyric_line["structure"]
|
||||
|
||||
# 2 优先选刚好包含在里面的
|
||||
full_spans = []
|
||||
partial_spans = []
|
||||
# print("non_structure_lyric_lines", non_structure_lyric_lines, n_chunks)
|
||||
for start_idx in range(len(valid_lyric_lines)):
|
||||
for end_idx in range(start_idx, len(valid_lyric_lines)):
|
||||
start = valid_lyric_lines[start_idx]["start"]
|
||||
end = start + cut_level * 10
|
||||
|
||||
# print("start_idx:", start_idx, "end_idx:", end_idx, "start:", start, "end:", end, "non_structure_lyric_lines[end_idx]:", non_structure_lyric_lines[end_idx])
|
||||
|
||||
if start_idx == end_idx and valid_lyric_lines[start_idx]["end"] > end:
|
||||
res = [(None, valid_lyric_lines[start_idx]["lyric_line"], valid_lyric_lines[start_idx]["phoneme_line_ipa"], (valid_lyric_lines[start_idx]["start"], valid_lyric_lines[start_idx]["end"])) for line_idx in range(start_idx, end_idx+1)]
|
||||
if len(res) > 0:
|
||||
partial_spans.append((start, end, res))
|
||||
break
|
||||
|
||||
if end_idx > 0 and end < valid_lyric_lines[end_idx]["start"] and valid_lyric_lines[end_idx-1]["end"] + start_pad_offset < end:
|
||||
res = [(None, valid_lyric_lines[line_idx]["lyric_line"], valid_lyric_lines[line_idx]["phoneme_line_ipa"], (valid_lyric_lines[line_idx]["start"], valid_lyric_lines[line_idx]["end"])) for line_idx in range(start_idx, end_idx)]
|
||||
if len(res) > 0:
|
||||
full_spans.append((start, end, res))
|
||||
break
|
||||
|
||||
if end < valid_lyric_lines[end_idx]["end"] + start_pad_offset and end > valid_lyric_lines[end_idx]["start"]:
|
||||
res = [(None, valid_lyric_lines[line_idx]["lyric_line"], valid_lyric_lines[line_idx]["phoneme_line_ipa"], (valid_lyric_lines[line_idx]["start"], valid_lyric_lines[line_idx]["end"])) for line_idx in range(start_idx, end_idx)]
|
||||
if len(res) > 0:
|
||||
partial_spans.append((start, end, res))
|
||||
break
|
||||
|
||||
if valid_lyric_lines[end_idx]["start"] > end:
|
||||
break
|
||||
|
||||
if start_idx == 0 and end_idx == len(valid_lyric_lines) - 1 and len(full_spans) == 0 and len(partial_spans) == 0:
|
||||
res = [(None, valid_lyric_lines[line_idx]["lyric_line"], valid_lyric_lines[line_idx]["phoneme_line_ipa"], (valid_lyric_lines[line_idx]["start"], valid_lyric_lines[line_idx]["end"])) for line_idx in range(start_idx, end_idx+1)]
|
||||
if len(res) > 0:
|
||||
full_spans.append((start, end, res))
|
||||
if expected_num_example is not None:
|
||||
if len(full_spans) >= expected_num_example or len(partial_spans) == 0:
|
||||
return full_spans
|
||||
if len(full_spans) + len(partial_spans) >= expected_num_example:
|
||||
left = expected_num_example - len(full_spans)
|
||||
return full_spans + random.sample(partial_spans, left)
|
||||
# print("full_spans:", full_spans)
|
||||
# print("partial_spans:", partial_spans)
|
||||
return full_spans + partial_spans
|
||||
|
||||
|
||||
class LyricProcessor(nn.Module):
|
||||
def __init__(self, infer=False):
|
||||
super().__init__()
|
||||
self.lyric_text_model = UMT5EncoderModel.from_pretrained("./checkpoints/umt5-base", local_files_only=True).eval().half()
|
||||
# not required gradient
|
||||
self.lyric_text_model.requires_grad_(False)
|
||||
self.lyric_text_tokenizer = AutoTokenizer.from_pretrained("./checkpoints/umt5-base", local_files_only=True)
|
||||
|
||||
|
||||
def get_text_embeddings(self, texts, device, text_max_length=256):
|
||||
inputs = self.lyric_text_tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=text_max_length)
|
||||
inputs = {key: value.to(device) for key, value in inputs.items()}
|
||||
if self.lyric_text_model.device != device:
|
||||
self.lyric_text_model.to(device)
|
||||
with torch.no_grad():
|
||||
outputs = self.lyric_text_model(**inputs)
|
||||
last_hidden_states = outputs.last_hidden_state
|
||||
attention_mask = inputs["attention_mask"]
|
||||
return last_hidden_states, attention_mask
|
||||
|
||||
def preprocess(self, valid_lyric_lines):
|
||||
lyric_texts = []
|
||||
ipa_texts = []
|
||||
for valid_line in valid_lyric_lines:
|
||||
structure, lyric_line, ipa_line = valid_line["structure"], valid_line["lyric"], valid_line["ipa"]
|
||||
if len(structure) > 0:
|
||||
lyric_texts.append(structure)
|
||||
if len(lyric_line) > 0:
|
||||
lyric_texts.append(lyric_line)
|
||||
if len(structure) == 0 and len(lyric_line) == 0:
|
||||
lyric_texts.append("")
|
||||
|
||||
if ipa_line != "_":
|
||||
ipa_line = self.split_unk(ipa_line.split(" "))
|
||||
ipa_line_str = " ".join(ipa_line)
|
||||
# 处理掉G2P的bug
|
||||
ipa_line_str = re.sub(r'\bz(?:\s+ə\s+z)+\b', "", ipa_line_str)
|
||||
ipa_line_str = re.sub(r'\s+', ' ', ipa_line_str).strip()
|
||||
ipa_texts.append(ipa_line_str)
|
||||
else:
|
||||
ipa_texts.append(ipa_line)
|
||||
|
||||
lyric_text = "\n".join(lyric_texts)
|
||||
ipa_text = " _ ".join(ipa_texts)
|
||||
return lyric_text, ipa_text
|
||||
|
||||
Reference in New Issue
Block a user