524 lines
16 KiB
Python
524 lines
16 KiB
Python
"""edge-tts backend for narration audio generation."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import re
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
from tts_backends.backend_common import publish_staged_pair, temporary_path
|
||
|
||
|
||
DEFAULT_SUBTITLE_MAX_CHARS = 20
|
||
DEFAULT_BOUNDARY_OVERLAP_TOLERANCE_MS = 100
|
||
_TICKS_PER_MILLISECOND = 10_000
|
||
_SENTENCE_END = frozenset("。！？!?")
|
||
_CLAUSE_END = frozenset("，,；;：:")
|
||
_CLOSING_PUNCTUATION = frozenset('”’」』）》)"\'')
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class _MappedWord:
|
||
start: int
|
||
end: int
|
||
source_start: int
|
||
source_end: int
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class _SubtitleCue:
|
||
start: int
|
||
end: int
|
||
text: str
|
||
|
||
|
||
COMMON_VOICES = [
|
||
("zh-CN", "zh-CN-XiaoxiaoNeural", "女声，普通话，清晰自然，默认推荐"),
|
||
("zh-CN", "zh-CN-XiaoyiNeural", "女声，普通话，明亮"),
|
||
("zh-CN", "zh-CN-YunjianNeural", "男声，普通话，稳重"),
|
||
("zh-CN", "zh-CN-YunxiNeural", "男声，普通话，年轻"),
|
||
("zh-CN", "zh-CN-YunxiaNeural", "男声，普通话，少年感"),
|
||
("zh-CN", "zh-CN-YunyangNeural", "男声，普通话，播报感"),
|
||
("zh-HK", "zh-HK-HiuGaaiNeural", "女声，粤语"),
|
||
("zh-HK", "zh-HK-WanLungNeural", "男声，粤语"),
|
||
("zh-TW", "zh-TW-HsiaoChenNeural", "女声，台湾普通话"),
|
||
("zh-TW", "zh-TW-YunJheNeural", "男声，台湾普通话"),
|
||
("en-US", "en-US-JennyNeural", "女声，美式英语"),
|
||
("en-US", "en-US-GuyNeural", "男声，美式英语"),
|
||
("en-GB", "en-GB-SoniaNeural", "女声，英式英语"),
|
||
("en-GB", "en-GB-RyanNeural", "男声，英式英语"),
|
||
]
|
||
|
||
|
||
def edge_output_extension() -> str:
|
||
return ".mp3"
|
||
|
||
|
||
def normalize_rate(rate: str) -> str:
|
||
"""Normalize a user-provided rate into edge-tts format."""
|
||
value = rate.strip()
|
||
if not value:
|
||
return "+0%"
|
||
if value.endswith("%"):
|
||
if value[0] not in "+-":
|
||
return f"+{value}"
|
||
return value
|
||
if re.fullmatch(r"[+-]?\d+", value):
|
||
number = int(value)
|
||
return f"{number:+d}%"
|
||
return value
|
||
|
||
|
||
async def generate(
|
||
text: str,
|
||
output_path: Path,
|
||
*,
|
||
voice: str,
|
||
rate: str,
|
||
subtitle_path: Path | None = None,
|
||
subtitle_max_chars: int = DEFAULT_SUBTITLE_MAX_CHARS,
|
||
) -> None:
|
||
"""Generate narration audio and, when requested, its compact SRT."""
|
||
if subtitle_path is not None:
|
||
await _generate_with_subtitles(
|
||
text,
|
||
output_path,
|
||
subtitle_path,
|
||
voice=voice,
|
||
rate=rate,
|
||
max_chars=subtitle_max_chars,
|
||
)
|
||
return
|
||
|
||
try:
|
||
import edge_tts
|
||
except ImportError as exc:
|
||
raise RuntimeError(
|
||
"Missing dependency `edge-tts`. Install it with: "
|
||
"python3 -m pip install edge-tts"
|
||
) from exc
|
||
|
||
communicate = edge_tts.Communicate(text, voice=voice, rate=normalize_rate(rate))
|
||
await communicate.save(str(output_path))
|
||
|
||
|
||
def _text_key(text: str) -> str:
|
||
return "".join(character.casefold() for character in text if character.isalnum())
|
||
|
||
|
||
def _source_key_positions(text: str) -> tuple[str, list[int]]:
|
||
key: list[str] = []
|
||
positions: list[int] = []
|
||
for index, character in enumerate(text):
|
||
if not character.isalnum():
|
||
continue
|
||
normalized = character.casefold()
|
||
key.extend(normalized)
|
||
positions.extend([index] * len(normalized))
|
||
return "".join(key), positions
|
||
|
||
|
||
def _map_word_boundaries(
|
||
text: str,
|
||
boundaries: list[dict],
|
||
provider_label: str,
|
||
) -> list[_MappedWord]:
|
||
source_key, source_positions = _source_key_positions(text)
|
||
boundary_keys = [_text_key(boundary["text"]) for boundary in boundaries]
|
||
boundary_key = "".join(boundary_keys)
|
||
if not source_key or source_key != boundary_key:
|
||
raise RuntimeError(
|
||
f"{provider_label} word boundaries could not be aligned with the narration text; "
|
||
"subtitle timing was not generated"
|
||
)
|
||
|
||
mapped: list[_MappedWord] = []
|
||
key_offset = 0
|
||
for boundary, word_key in zip(boundaries, boundary_keys):
|
||
if not word_key:
|
||
continue
|
||
key_end = key_offset + len(word_key)
|
||
mapped.append(
|
||
_MappedWord(
|
||
start=boundary["offset"],
|
||
end=boundary["offset"] + boundary["duration"],
|
||
source_start=source_positions[key_offset],
|
||
source_end=source_positions[key_end - 1] + 1,
|
||
)
|
||
)
|
||
key_offset = key_end
|
||
return mapped
|
||
|
||
|
||
def _trim_span(text: str, start: int, end: int) -> tuple[int, int]:
|
||
while start < end and text[start].isspace():
|
||
start += 1
|
||
while end > start and text[end - 1].isspace():
|
||
end -= 1
|
||
return start, end
|
||
|
||
|
||
def _display_length(text: str, start: int, end: int) -> int:
|
||
return sum(not character.isspace() for character in text[start:end])
|
||
|
||
|
||
def _is_sentence_end(text: str, index: int) -> bool:
|
||
character = text[index]
|
||
if character in _SENTENCE_END:
|
||
return True
|
||
if character != ".":
|
||
return False
|
||
previous = text[index - 1] if index else ""
|
||
following = text[index + 1] if index + 1 < len(text) else ""
|
||
return not (previous.isdigit() and following.isdigit())
|
||
|
||
|
||
def _sentence_spans(text: str) -> list[tuple[int, int]]:
|
||
spans: list[tuple[int, int]] = []
|
||
start = 0
|
||
index = 0
|
||
while index < len(text):
|
||
if not _is_sentence_end(text, index):
|
||
index += 1
|
||
continue
|
||
end = index + 1
|
||
while end < len(text) and text[end] in _CLOSING_PUNCTUATION:
|
||
end += 1
|
||
span = _trim_span(text, start, end)
|
||
if span[0] < span[1]:
|
||
spans.append(span)
|
||
start = end
|
||
index = end
|
||
span = _trim_span(text, start, len(text))
|
||
if span[0] < span[1]:
|
||
spans.append(span)
|
||
return spans
|
||
|
||
|
||
def _hard_split_span(
|
||
text: str,
|
||
span: tuple[int, int],
|
||
words: list[_MappedWord],
|
||
max_chars: int,
|
||
) -> list[tuple[int, int]]:
|
||
start, end = span
|
||
parts: list[tuple[int, int]] = []
|
||
while _display_length(text, start, end) > max_chars:
|
||
remaining_length = _display_length(text, start, end)
|
||
remaining_parts = (remaining_length + max_chars - 1) // max_chars
|
||
target_length = (remaining_length + remaining_parts - 1) // remaining_parts
|
||
candidates = [
|
||
(word.source_end, _display_length(text, start, word.source_end))
|
||
for word in words
|
||
if start < word.source_end < end
|
||
and _display_length(text, start, word.source_end) <= max_chars
|
||
]
|
||
if candidates:
|
||
split_at, _ = min(
|
||
candidates,
|
||
key=lambda candidate: (
|
||
abs(candidate[1] - target_length),
|
||
-candidate[1],
|
||
),
|
||
)
|
||
else:
|
||
split_at = next(
|
||
(
|
||
word.source_end
|
||
for word in words
|
||
if start < word.source_end < end
|
||
),
|
||
end,
|
||
)
|
||
if split_at >= end:
|
||
break
|
||
part = _trim_span(text, start, split_at)
|
||
if part[0] < part[1]:
|
||
parts.append(part)
|
||
start = split_at
|
||
part = _trim_span(text, start, end)
|
||
if part[0] < part[1]:
|
||
parts.append(part)
|
||
return parts
|
||
|
||
|
||
def _split_sentence_span(
|
||
text: str,
|
||
sentence: tuple[int, int],
|
||
words: list[_MappedWord],
|
||
max_chars: int,
|
||
) -> list[tuple[int, int]]:
|
||
if _display_length(text, *sentence) <= max_chars:
|
||
return [sentence]
|
||
|
||
start, end = sentence
|
||
clauses: list[tuple[int, int]] = []
|
||
clause_start = start
|
||
for index in range(start, end):
|
||
if text[index] not in _CLAUSE_END:
|
||
continue
|
||
clause = _trim_span(text, clause_start, index + 1)
|
||
if clause[0] < clause[1]:
|
||
clauses.append(clause)
|
||
clause_start = index + 1
|
||
clause = _trim_span(text, clause_start, end)
|
||
if clause[0] < clause[1]:
|
||
clauses.append(clause)
|
||
|
||
atoms = [
|
||
part
|
||
for clause in clauses
|
||
for part in _hard_split_span(text, clause, words, max_chars)
|
||
]
|
||
merged: list[tuple[int, int]] = []
|
||
for atom in atoms:
|
||
if not merged:
|
||
merged.append(atom)
|
||
continue
|
||
candidate = (merged[-1][0], atom[1])
|
||
if _display_length(text, *candidate) <= max_chars:
|
||
merged[-1] = candidate
|
||
else:
|
||
merged.append(atom)
|
||
return merged
|
||
|
||
|
||
def _clamp_small_overlaps(
|
||
cues: list[_SubtitleCue],
|
||
*,
|
||
tolerance_ms: int = DEFAULT_BOUNDARY_OVERLAP_TOLERANCE_MS,
|
||
provider_label: str,
|
||
) -> list[_SubtitleCue]:
|
||
tolerance = tolerance_ms * _TICKS_PER_MILLISECOND
|
||
normalized: list[_SubtitleCue] = []
|
||
for cue in cues:
|
||
if normalized and cue.start < normalized[-1].end:
|
||
overlap = normalized[-1].end - cue.start
|
||
if overlap > tolerance:
|
||
overlap_ms = overlap / _TICKS_PER_MILLISECOND
|
||
raise RuntimeError(
|
||
f"{provider_label} returned {overlap_ms:g} ms of overlapping "
|
||
"word-boundary timing; audio and subtitles were not published"
|
||
)
|
||
cue = _SubtitleCue(
|
||
start=normalized[-1].end,
|
||
end=cue.end,
|
||
text=cue.text,
|
||
)
|
||
if cue.end <= cue.start:
|
||
raise RuntimeError(
|
||
f"{provider_label} returned an invalid subtitle timing interval; "
|
||
"audio and subtitles were not published"
|
||
)
|
||
normalized.append(cue)
|
||
return normalized
|
||
|
||
|
||
def _subtitle_cues(
|
||
text: str,
|
||
boundaries: list[dict],
|
||
max_chars: int,
|
||
provider_label: str,
|
||
) -> list[_SubtitleCue]:
|
||
if max_chars < 1:
|
||
raise ValueError("subtitle_max_chars must be at least 1")
|
||
words = _map_word_boundaries(text, boundaries, provider_label)
|
||
spans = [
|
||
span
|
||
for sentence in _sentence_spans(text)
|
||
for span in _split_sentence_span(text, sentence, words, max_chars)
|
||
]
|
||
|
||
pending: list[tuple[int, int, str]] = []
|
||
assigned_word_indexes: list[int] = []
|
||
for start, end in spans:
|
||
matching = [
|
||
(index, word)
|
||
for index, word in enumerate(words)
|
||
if word.source_start >= start and word.source_end <= end
|
||
]
|
||
if not matching:
|
||
continue
|
||
cue_text = re.sub(r"\s+", " ", text[start:end]).strip()
|
||
assigned_word_indexes.extend(index for index, _ in matching)
|
||
pending.append((matching[0][1].start, matching[-1][1].end, cue_text))
|
||
|
||
if assigned_word_indexes != list(range(len(words))):
|
||
raise RuntimeError(
|
||
f"{provider_label} word boundaries crossed subtitle split points; "
|
||
"subtitle timing was not generated"
|
||
)
|
||
if not pending:
|
||
raise RuntimeError(f"{provider_label} produced no timed subtitle cues")
|
||
|
||
cues: list[_SubtitleCue] = []
|
||
for index, (start, word_end, cue_text) in enumerate(pending):
|
||
next_start = pending[index + 1][0] if index + 1 < len(pending) else None
|
||
end = next_start if next_start is not None and next_start > word_end else word_end
|
||
cues.append(_SubtitleCue(start=start, end=end, text=cue_text))
|
||
cues = _clamp_small_overlaps(cues, provider_label=provider_label)
|
||
|
||
source_text = re.sub(r"\s+", "", text)
|
||
subtitle_text = re.sub(r"\s+", "", "".join(cue.text for cue in cues))
|
||
if subtitle_text != source_text:
|
||
raise RuntimeError(
|
||
f"{provider_label} subtitle text does not match the narration text; "
|
||
"audio and subtitles were not published"
|
||
)
|
||
if any(_display_length(cue.text, 0, len(cue.text)) > max_chars for cue in cues):
|
||
raise RuntimeError(
|
||
f"A single {provider_label} word boundary exceeds the subtitle character limit; "
|
||
"audio and subtitles were not published"
|
||
)
|
||
return cues
|
||
|
||
|
||
def _srt_timestamp(ticks: int) -> str:
|
||
total_milliseconds = round(ticks / 10_000)
|
||
hours, remainder = divmod(total_milliseconds, 3_600_000)
|
||
minutes, remainder = divmod(remainder, 60_000)
|
||
seconds, milliseconds = divmod(remainder, 1_000)
|
||
return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"
|
||
|
||
|
||
def _format_srt(cues: list[_SubtitleCue]) -> str:
|
||
blocks = [
|
||
(
|
||
f"{index}\n"
|
||
f"{_srt_timestamp(cue.start)} --> {_srt_timestamp(cue.end)}\n"
|
||
f"{cue.text}"
|
||
)
|
||
for index, cue in enumerate(cues, 1)
|
||
]
|
||
return "\n\n".join(blocks) + "\n"
|
||
|
||
|
||
def format_word_timed_srt(
|
||
text: str,
|
||
boundaries: list[dict],
|
||
max_chars: int = DEFAULT_SUBTITLE_MAX_CHARS,
|
||
*,
|
||
provider_label: str = "Edge TTS",
|
||
) -> str:
|
||
"""Regroup provider word timings into compact, text-faithful SRT cues."""
|
||
return _format_srt(
|
||
_subtitle_cues(
|
||
text,
|
||
boundaries,
|
||
max_chars,
|
||
provider_label,
|
||
)
|
||
)
|
||
|
||
|
||
async def _generate_with_subtitles(
|
||
text: str,
|
||
output_path: Path,
|
||
subtitle_path: Path,
|
||
*,
|
||
voice: str,
|
||
rate: str,
|
||
max_chars: int,
|
||
) -> None:
|
||
"""Generate one MP3 and compact SRT from the same Edge word-timing stream."""
|
||
try:
|
||
import edge_tts
|
||
except ImportError as exc:
|
||
raise RuntimeError(
|
||
"Missing dependency `edge-tts`. Install it with: "
|
||
"python3 -m pip install edge-tts"
|
||
) from exc
|
||
|
||
communicate = edge_tts.Communicate(
|
||
text,
|
||
voice=voice,
|
||
rate=normalize_rate(rate),
|
||
boundary="WordBoundary",
|
||
)
|
||
audio_descriptor = -1
|
||
subtitle_descriptor = -1
|
||
staged_audio: Path | None = None
|
||
staged_subtitle: Path | None = None
|
||
boundaries: list[dict] = []
|
||
received_audio = False
|
||
try:
|
||
audio_descriptor, staged_audio = temporary_path(output_path, ".tmp")
|
||
subtitle_descriptor, staged_subtitle = temporary_path(subtitle_path, ".tmp")
|
||
|
||
audio_stream = os.fdopen(audio_descriptor, "wb")
|
||
audio_descriptor = -1
|
||
with audio_stream:
|
||
async for chunk in communicate.stream():
|
||
if chunk["type"] == "audio":
|
||
audio_stream.write(chunk["data"])
|
||
received_audio = True
|
||
elif chunk["type"] == "WordBoundary":
|
||
boundaries.append(chunk)
|
||
audio_stream.flush()
|
||
os.fsync(audio_stream.fileno())
|
||
|
||
if not received_audio:
|
||
raise RuntimeError("Edge TTS returned no audio data")
|
||
if not boundaries:
|
||
raise RuntimeError("Edge TTS returned no word-boundary timing")
|
||
subtitle_text = format_word_timed_srt(text, boundaries, max_chars)
|
||
|
||
subtitle_stream = os.fdopen(
|
||
subtitle_descriptor,
|
||
"w",
|
||
encoding="utf-8",
|
||
newline="\n",
|
||
)
|
||
subtitle_descriptor = -1
|
||
with subtitle_stream:
|
||
subtitle_stream.write(subtitle_text)
|
||
subtitle_stream.flush()
|
||
os.fsync(subtitle_stream.fileno())
|
||
assert staged_audio is not None
|
||
assert staged_subtitle is not None
|
||
publish_staged_pair(
|
||
staged_audio,
|
||
output_path,
|
||
staged_subtitle,
|
||
subtitle_path,
|
||
)
|
||
finally:
|
||
if audio_descriptor >= 0:
|
||
os.close(audio_descriptor)
|
||
if subtitle_descriptor >= 0:
|
||
os.close(subtitle_descriptor)
|
||
if staged_audio is not None:
|
||
staged_audio.unlink(missing_ok=True)
|
||
if staged_subtitle is not None:
|
||
staged_subtitle.unlink(missing_ok=True)
|
||
|
||
|
||
def print_common_voices() -> None:
|
||
print("Common edge-tts voices:")
|
||
print("Locale Voice Notes")
|
||
print("------ ---------------------------- ----------------")
|
||
for locale, voice, notes in COMMON_VOICES:
|
||
print(f"{locale:<8} {voice:<29} {notes}")
|
||
|
||
|
||
async def print_voices(locale: str | None = None) -> None:
|
||
try:
|
||
import edge_tts
|
||
except ImportError as exc:
|
||
raise RuntimeError(
|
||
"Missing dependency `edge-tts`. Install it with: "
|
||
"python3 -m pip install edge-tts"
|
||
) from exc
|
||
|
||
manager = await edge_tts.VoicesManager.create()
|
||
voices = manager.voices
|
||
if locale:
|
||
voices = [voice for voice in voices if voice.get("Locale") == locale]
|
||
for voice in sorted(voices, key=lambda item: (item.get("Locale", ""), item.get("ShortName", ""))):
|
||
short_name = voice.get("ShortName", "")
|
||
voice_locale = voice.get("Locale", "")
|
||
gender = voice.get("Gender", "")
|
||
friendly = voice.get("FriendlyName", "")
|
||
print(f"{voice_locale:<8} {short_name:<34} {gender:<8} {friendly}")
|