Mirek Sobczak
279 lines
8.2 KiB
Python
279 lines
8.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert a long text file to MP3 using the existing Riva TTS client.
|
|
|
|
The script splits input text into smaller chunks, synthesizes each chunk with
|
|
`python-clients/scripts/tts/talk.py`, joins all WAV chunks, and finally converts
|
|
the result to MP3 using ffmpeg.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import wave
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Convert a long TXT file to MP3 using NVIDIA Riva TTS.",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
|
)
|
|
parser.add_argument("--input", type=Path, required=True, help="Path to input .txt file.")
|
|
parser.add_argument("--output", type=Path, required=True, help="Path to output .mp3 file.")
|
|
parser.add_argument("--server", default="0.0.0.0:50051", help="Riva server address.")
|
|
parser.add_argument("--language-code", default="pl-PL", help="Language code for TTS.")
|
|
parser.add_argument(
|
|
"--voice",
|
|
default="Chatterbox-Multilingual.pl-PL.Male",
|
|
help="Riva voice name.",
|
|
)
|
|
parser.add_argument(
|
|
"--talk-script",
|
|
type=Path,
|
|
default=Path("python-clients/scripts/tts/talk.py"),
|
|
help="Path to existing talk.py script.",
|
|
)
|
|
parser.add_argument(
|
|
"--max-chars",
|
|
type=int,
|
|
default=700,
|
|
help="Maximum characters per TTS chunk.",
|
|
)
|
|
parser.add_argument(
|
|
"--sample-rate-hz",
|
|
type=int,
|
|
default=22050,
|
|
help="Output sample rate used by talk.py.",
|
|
)
|
|
parser.add_argument(
|
|
"--keep-wav",
|
|
action="store_true",
|
|
help="Keep final merged WAV next to MP3.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def normalize_whitespace(text: str) -> str:
|
|
text = text.replace("\r\n", "\n").replace("\r", "\n")
|
|
lines = [line.strip() for line in text.split("\n")]
|
|
return "\n".join(lines)
|
|
|
|
|
|
def split_sentence_if_needed(sentence: str, max_chars: int) -> List[str]:
|
|
if len(sentence) <= max_chars:
|
|
return [sentence]
|
|
|
|
words = sentence.split()
|
|
parts: List[str] = []
|
|
current = ""
|
|
for word in words:
|
|
candidate = f"{current} {word}".strip() if current else word
|
|
if len(candidate) <= max_chars:
|
|
current = candidate
|
|
else:
|
|
if current:
|
|
parts.append(current)
|
|
# Single very long token fallback.
|
|
if len(word) > max_chars:
|
|
for i in range(0, len(word), max_chars):
|
|
parts.append(word[i : i + max_chars])
|
|
current = ""
|
|
else:
|
|
current = word
|
|
if current:
|
|
parts.append(current)
|
|
return parts
|
|
|
|
|
|
def chunk_text(text: str, max_chars: int) -> List[str]:
|
|
paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
|
|
sentence_splitter = re.compile(r"(?<=[.!?…])\s+")
|
|
|
|
chunks: List[str] = []
|
|
current = ""
|
|
|
|
for paragraph in paragraphs:
|
|
sentences = [s.strip() for s in sentence_splitter.split(paragraph) if s.strip()]
|
|
expanded_sentences: List[str] = []
|
|
for s in sentences:
|
|
expanded_sentences.extend(split_sentence_if_needed(s, max_chars))
|
|
|
|
for sentence in expanded_sentences:
|
|
candidate = f"{current} {sentence}".strip() if current else sentence
|
|
if len(candidate) <= max_chars:
|
|
current = candidate
|
|
else:
|
|
if current:
|
|
chunks.append(current)
|
|
current = sentence
|
|
|
|
if current:
|
|
chunks.append(current)
|
|
current = ""
|
|
|
|
if current:
|
|
chunks.append(current)
|
|
|
|
if not chunks:
|
|
raise ValueError("Input text is empty after preprocessing.")
|
|
|
|
return chunks
|
|
|
|
|
|
def synthesize_chunk(
|
|
talk_script: Path,
|
|
server: str,
|
|
language_code: str,
|
|
voice: str,
|
|
sample_rate_hz: int,
|
|
text: str,
|
|
output_wav: Path,
|
|
) -> None:
|
|
cmd = [
|
|
sys.executable,
|
|
str(talk_script),
|
|
"--server",
|
|
server,
|
|
"--language-code",
|
|
language_code,
|
|
"--voice",
|
|
voice,
|
|
"--text",
|
|
text,
|
|
"--stream",
|
|
"--sample-rate-hz",
|
|
str(sample_rate_hz),
|
|
"--output",
|
|
str(output_wav),
|
|
]
|
|
|
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(
|
|
"TTS synthesis failed.\n"
|
|
f"Command: {' '.join(cmd)}\n"
|
|
f"stdout:\n{result.stdout}\n"
|
|
f"stderr:\n{result.stderr}"
|
|
)
|
|
|
|
|
|
def merge_wav_files(wav_files: List[Path], output_wav: Path) -> None:
|
|
if not wav_files:
|
|
raise ValueError("No WAV files to merge.")
|
|
|
|
with wave.open(str(wav_files[0]), "rb") as first:
|
|
params = first.getparams()
|
|
frames = [first.readframes(first.getnframes())]
|
|
|
|
for wav_path in wav_files[1:]:
|
|
with wave.open(str(wav_path), "rb") as wf:
|
|
if (
|
|
wf.getnchannels() != params.nchannels
|
|
or wf.getsampwidth() != params.sampwidth
|
|
or wf.getframerate() != params.framerate
|
|
or wf.getcomptype() != params.comptype
|
|
):
|
|
raise RuntimeError(f"Incompatible WAV params in {wav_path}")
|
|
frames.append(wf.readframes(wf.getnframes()))
|
|
|
|
with wave.open(str(output_wav), "wb") as out:
|
|
out.setparams(params)
|
|
for chunk in frames:
|
|
out.writeframes(chunk)
|
|
|
|
|
|
def convert_wav_to_mp3(input_wav: Path, output_mp3: Path) -> None:
|
|
ffmpeg = shutil.which("ffmpeg")
|
|
if ffmpeg is None:
|
|
raise RuntimeError(
|
|
"ffmpeg not found in PATH. Install ffmpeg to generate MP3 output."
|
|
)
|
|
|
|
cmd = [
|
|
ffmpeg,
|
|
"-y",
|
|
"-i",
|
|
str(input_wav),
|
|
"-codec:a",
|
|
"libmp3lame",
|
|
"-q:a",
|
|
"2",
|
|
str(output_mp3),
|
|
]
|
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(
|
|
"MP3 conversion failed.\n"
|
|
f"Command: {' '.join(cmd)}\n"
|
|
f"stdout:\n{result.stdout}\n"
|
|
f"stderr:\n{result.stderr}"
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
|
|
if args.max_chars < 50:
|
|
raise ValueError("--max-chars should be at least 50.")
|
|
|
|
input_path = args.input.expanduser().resolve()
|
|
output_mp3 = args.output.expanduser().resolve()
|
|
talk_script = args.talk_script.expanduser().resolve()
|
|
|
|
if not input_path.is_file():
|
|
raise FileNotFoundError(f"Input file not found: {input_path}")
|
|
if not talk_script.is_file():
|
|
raise FileNotFoundError(f"talk.py not found: {talk_script}")
|
|
|
|
output_mp3.parent.mkdir(parents=True, exist_ok=True)
|
|
merged_wav = output_mp3.with_suffix(".wav")
|
|
|
|
text = input_path.read_text(encoding="utf-8")
|
|
text = normalize_whitespace(text)
|
|
chunks = chunk_text(text, args.max_chars)
|
|
|
|
print(f"Input split into {len(chunks)} chunks (max {args.max_chars} chars each).")
|
|
|
|
with tempfile.TemporaryDirectory(prefix="tts_chunks_") as tmp_dir:
|
|
tmp_path = Path(tmp_dir)
|
|
wav_paths: List[Path] = []
|
|
|
|
for idx, chunk in enumerate(chunks, start=1):
|
|
chunk_wav = tmp_path / f"chunk_{idx:05d}.wav"
|
|
print(f"[{idx}/{len(chunks)}] Synthesizing chunk...")
|
|
synthesize_chunk(
|
|
talk_script=talk_script,
|
|
server=args.server,
|
|
language_code=args.language_code,
|
|
voice=args.voice,
|
|
sample_rate_hz=args.sample_rate_hz,
|
|
text=chunk,
|
|
output_wav=chunk_wav,
|
|
)
|
|
wav_paths.append(chunk_wav)
|
|
|
|
print("Merging WAV chunks...")
|
|
merge_wav_files(wav_paths, merged_wav)
|
|
|
|
print("Converting WAV to MP3...")
|
|
convert_wav_to_mp3(merged_wav, output_mp3)
|
|
|
|
if not args.keep_wav and merged_wav.exists():
|
|
merged_wav.unlink()
|
|
|
|
print(f"Done: {output_mp3}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
raise SystemExit(main())
|
|
except Exception as exc: # pragma: no cover
|
|
print(f"Error: {exc}", file=sys.stderr)
|
|
raise SystemExit(1) |