135 lines
4.3 KiB
Python
135 lines
4.3 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
import io
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
from loguru import logger
|
|
from pydantic import BaseModel
|
|
from kokoro import KPipeline
|
|
import soundfile as sf
|
|
|
|
# --- Configuration ---
|
|
DATA_DIR = os.environ.get("DATA_DIR", "data")
|
|
INPUT_FILE = os.path.join(DATA_DIR, "scripts.json")
|
|
OUTPUT_DIR = Path(DATA_DIR) / "audio"
|
|
VOICE = os.environ.get("TTS_VOICE", "af_heart")
|
|
CONCURRENCY_LIMIT = 1
|
|
|
|
class ScriptItem(BaseModel):
|
|
headline: str
|
|
script: str
|
|
|
|
class KokoroTTS:
|
|
def __init__(self, lang_code: str = 'a'):
|
|
self.pipeline = KPipeline(lang_code=lang_code)
|
|
|
|
async def text_to_speech(self, text: str, output_wav: Path):
|
|
"""
|
|
Generates speech locally using Kokoro and saves the resulting audio to a WAV file.
|
|
"""
|
|
try:
|
|
logger.debug(f"Generating TTS locally for: {text[:50]}...")
|
|
|
|
# The pipeline returns a generator
|
|
generator = self.pipeline(text, voice=VOICE)
|
|
|
|
# We take the first audio segment produced
|
|
# For short scripts, it's usually one segment.
|
|
# For longer ones, we concatenate.
|
|
all_audio = []
|
|
for _, _, audio in generator:
|
|
all_audio.append(audio)
|
|
|
|
if not all_audio:
|
|
raise RuntimeError("No audio data generated by Kokoro")
|
|
|
|
import numpy as np
|
|
final_audio = np.concatenate(all_audio)
|
|
|
|
# Kokoro outputs at 24000 Hz
|
|
sf.write(output_wav, final_audio, 24000)
|
|
|
|
except Exception as e:
|
|
logger.exception(f"Unexpected error during Kokoro TTS: {e}")
|
|
raise
|
|
|
|
async def convert_wav_to_mp3(wav_path: Path, mp3_path: Path):
|
|
"""
|
|
Converts a WAV file to MP3 using ffmpeg.
|
|
"""
|
|
try:
|
|
cmd = ["ffmpeg", "-y", "-i", str(wav_path), "-q:a", "2", str(mp3_path)]
|
|
process = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE
|
|
)
|
|
stdout, stderr = await process.communicate()
|
|
|
|
if process.returncode != 0:
|
|
logger.error(f"ffmpeg error: {stderr.decode()}")
|
|
return False
|
|
return True
|
|
except Exception as e:
|
|
logger.exception(f"Unexpected error during conversion: {e}")
|
|
return False
|
|
|
|
async def process_article(semaphore: asyncio.Semaphore, client: KokoroTTS, item: ScriptItem):
|
|
"""
|
|
Processes a single article: combines text, generates audio, and converts to MP3.
|
|
"""
|
|
async with semaphore:
|
|
if not item.headline:
|
|
logger.warning("Skipping article processing: headline is empty.")
|
|
return
|
|
safe_headline = "".join([c if c.isalnum() else "_" for c in item.headline])[:100]
|
|
wav_path = OUTPUT_DIR / f"{safe_headline}.wav"
|
|
mp3_path = OUTPUT_DIR / f"{safe_headline}.mp3"
|
|
|
|
full_text = f"{item.headline}. {item.script}"
|
|
|
|
try:
|
|
await client.text_to_speech(full_text, wav_path)
|
|
|
|
if await convert_wav_to_mp3(wav_path, mp3_path):
|
|
logger.info(f"Successfully created audio for: {item.headline}")
|
|
else:
|
|
logger.error(f"Failed to convert audio for: {item.headline}")
|
|
|
|
if wav_path.exists():
|
|
wav_path.unlink()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing '{item.headline}': {e}")
|
|
|
|
async def main():
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
try:
|
|
with open(INPUT_FILE, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
scripts = [ScriptItem(**item) for item in data]
|
|
except Exception as e:
|
|
logger.error(f"Failed to load input file {INPUT_FILE}: {e}")
|
|
return
|
|
|
|
client = KokoroTTS()
|
|
|
|
logger.info(f"Loaded {len(scripts)} scripts. Starting local TTS generation via Kokoro...")
|
|
semaphore = asyncio.Semaphore(CONCURRENCY_LIMIT)
|
|
|
|
tasks = [process_article(semaphore, client, item) for item in scripts]
|
|
|
|
total = len(scripts)
|
|
for i, coro in enumerate(asyncio.as_completed(tasks), 1):
|
|
await coro
|
|
logger.info(f"🚀 Progress: 🎙️ {i}/{total} articles processed! ✨")
|
|
|
|
logger.info("TTS generation complete.")
|
|
|
|
if __name__ == "__main__":
|
|
logger.add("tts_generator.log", rotation="1 MB")
|
|
asyncio.run(main())
|