initial commit

tts_generator.py

@@ -0,0 +1,133 @@
+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 ---
+INPUT_FILE = "data/scripts.json"
+OUTPUT_DIR = Path("data/audio")
+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())