satoko-chatbot/voice_tts.py

"""Synthèse HTTP interchangeable et lecture Discord, sans fichiers temporaires."""
import asyncio
import audioop
import logging
import os
import re
import time
import wave
from io import BytesIO

import discord
import emoji
import httpx

logger = logging.getLogger(__name__)


def speech_chunks(text, limit=240, first_limit=None):
    if limit < 1 or (first_limit is not None and first_limit < 1):
        raise ValueError('La taille des morceaux doit être positive')
    first = True
    text = re.sub(r'```.*?```', ' Bloc de code omis. ', text, flags=re.S)
    text = re.sub(r'\[([^\]]+)\]\([^)]+\)', r'\1', text)
    text = re.sub(r'https?://\S+|<a?:\w+:\d+>|<[@#][!&]?\d+>', ' ', text)
    # Le prompt peut émettre des alias personnalisés (:hap:, :desu:, etc.).
    # Exiger une lettre évite de confondre ces alias avec une heure 12:30:45.
    text = re.sub(r':[A-Za-z_][A-Za-z0-9_+\-]*:', ' ', text)
    text = emoji.replace_emoji(text, replace=' ')
    text = re.sub(r'[*_`~#>|]', '', text)
    text = re.sub(r'[^\S\n]+', ' ', text)
    for sentence in re.split(r'(?<=[.!?…])\s+|\n+', text.strip()):
        while sentence:
            budget = min(limit, first_limit) if first and first_limit else limit
            end = len(sentence) if len(sentence) <= budget else sentence.rfind(' ', 0, budget + 1)
            # Favoriser une pause naturelle, sans fabriquer un minuscule morceau.
            if len(sentence) > budget:
                pauses = list(re.finditer(r'[,;:]\s', sentence[:budget + 1]))
                if pauses and pauses[-1].start() >= budget // 2:
                    end = pauses[-1].start() + 1
            if end <= 0:
                end = budget
            chunk, sentence = sentence[:end].strip(), sentence[end:].strip()
            if chunk:
                first = False
                yield chunk


class WavSource(discord.AudioSource):
    def __init__(self, payload):
        with wave.open(BytesIO(payload), 'rb') as wav:
            if wav.getsampwidth() != 2 or wav.getnchannels() not in (1, 2):
                raise ValueError('Le TTS doit retourner du WAV PCM 16 bits mono/stéréo')
            pcm = wav.readframes(wav.getnframes())
            if wav.getnchannels() == 1:
                pcm = audioop.tostereo(pcm, 2, 1, 1)
            pcm, _ = audioop.ratecv(pcm, 2, 2, wav.getframerate(), 48000, None)
        self.audio = BytesIO(pcm)

    def read(self):
        frame = self.audio.read(3840)  # Discord : 20 ms, 48 kHz, stéréo, PCM16.
        return frame.ljust(3840, b'\0') if frame else b''

    def cleanup(self):
        self.audio.close()


async def play_wav(voice, payload):
    loop = asyncio.get_running_loop()
    done = loop.create_future()
    source = WavSource(payload)

    def complete(error):
        if not done.done():
            if error:
                done.set_exception(error)
            else:
                done.set_result(None)

    try:
        voice.play(source, after=lambda error: loop.call_soon_threadsafe(complete, error))
        await done
    finally:
        if not done.done() or done.cancelled():
            voice.stop()
        source.cleanup()


async def speak(voice, text, url, voice_name='ff_siwis'):
    first_chars = int(os.getenv('TTS_FIRST_CHUNK_CHARS', '48'))
    chunk_chars = int(os.getenv('TTS_CHUNK_CHARS', '80'))
    if not (16 <= first_chars <= chunk_chars <= 240):
        raise ValueError('Tailles TTS attendues : 16 <= premier <= suivants <= 240')
    chunks = iter(speech_chunks(text, limit=chunk_chars, first_limit=first_chars))
    started = time.monotonic()
    first_playback = True
    async with httpx.AsyncClient(timeout=60) as http:
        async def synthesize(chunk):
            async with http.stream('POST', url, json={'input': chunk, 'voice': voice_name}) as response:
                response.raise_for_status()
                payload = bytearray()
                async for data in response.aiter_bytes():
                    payload.extend(data)
                    if len(payload) > 8 * 1024 * 1024:
                        raise ValueError('Réponse audio trop volumineuse')
                return bytes(payload)

        first = next(chunks, None)
        if first is None:
            return  # Une réponse uniquement composée d'emojis reste silencieuse.
        pending = asyncio.create_task(synthesize(first))
        try:
            while pending is not None:
                payload = await pending
                following = next(chunks, None)
                # Une seule phrase préchargée pendant la lecture de la précédente.
                pending = asyncio.create_task(synthesize(following)) if following else None
                if first_playback:
                    logger.info('Premier audio prêt : délai TTS=%.3fs', time.monotonic() - started)
                    first_playback = False
                await play_wav(voice, payload)
        finally:
            if pending is not None:
                pending.cancel()
                await asyncio.gather(pending, return_exceptions=True)