dnnai/tts-qt-client/tts_server/server.py

#!/usr/bin/env python3
import asyncio
import base64
import json
import argparse
import signal
import sys
import os
import io
import wave
import hashlib
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache
from typing import Set, Optional, Tuple, List, Dict
import numpy as np
import websockets
from websockets.server import WebSocketServerProtocol

sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)


class AudioCache:
    def __init__(self, max_size: int = 200):
        self.cache: Dict[str, bytes] = {}
        self.access_order: List[str] = []
        self.max_size = max_size
        self.hits = 0
        self.misses = 0
    
    def _make_key(self, text: str, ref_hash: str) -> str:
        return hashlib.md5(f"{ref_hash}:{text}".encode()).hexdigest()
    
    def get(self, text: str, ref_hash: str) -> Optional[bytes]:
        key = self._make_key(text, ref_hash)
        if key in self.cache:
            self.access_order.remove(key)
            self.access_order.append(key)
            self.hits += 1
            return self.cache[key]
        self.misses += 1
        return None
    
    def put(self, text: str, ref_hash: str, audio_data: bytes):
        key = self._make_key(text, ref_hash)
        if key in self.cache:
            self.access_order.remove(key)
        elif len(self.cache) >= self.max_size:
            oldest = self.access_order.pop(0)
            del self.cache[oldest]
        self.cache[key] = audio_data
        self.access_order.append(key)
    
    def get_stats(self) -> dict:
        total = self.hits + self.misses
        hit_rate = (self.hits / total * 100) if total > 0 else 0
        return {
            'size': len(self.cache),
            'max_size': self.max_size,
            'hits': self.hits,
            'misses': self.misses,
            'hit_rate': f"{hit_rate:.1f}%"
        }


class TTSEngine:
    def __init__(self, model_id: str = "Qwen/Qwen3-TTS-12Hz-0.6B-Base", 
                 ref_audio_path: str = None):
        import soundfile as sf
        
        self.model_id = model_id
        self.sample_rate = 24000
        self.ref_audio_path = ref_audio_path
        self.ref_hash = None
        self.audio_cache = AudioCache(max_size=200)
        
        print(f"Loading Qwen3-TTS model: {model_id}")
        from qwen_tts import Qwen3TTSModel
        self.model = Qwen3TTSModel.from_pretrained(model_id, device_map='mps')
        print("Model loaded successfully")
        
        if self.ref_audio_path and os.path.exists(self.ref_audio_path):
            self.ref_audio, self.ref_sr = self._load_audio(self.ref_audio_path)
            self.ref_hash = hashlib.md5(open(self.ref_audio_path, 'rb').read()).hexdigest()
            self.prompt_items = self.model.create_voice_clone_prompt(
                ref_audio=(self.ref_audio, self.ref_sr),
                ref_text="",
                x_vector_only_mode=True
            )
            print(f"Reference audio loaded: {self.ref_audio_path}")
        else:
            self.ref_audio = None
            self.ref_sr = None
            self.prompt_items = None
            print("Warning: No reference audio provided")
    
    def _load_audio(self, path: str) -> Tuple[np.ndarray, int]:
        import soundfile as sf
        audio, sr = sf.read(path)
        if len(audio.shape) > 1:
            audio = audio[:, 0]
        return audio, sr
    
    def synthesize(self, text: str, ref_audio_path: str = None) -> bytes:
        if ref_audio_path:
            ref_audio, ref_sr = self._load_audio(ref_audio_path)
            ref_hash = hashlib.md5(open(ref_audio_path, 'rb').read()).hexdigest()
            prompt = self.model.create_voice_clone_prompt(
                ref_audio=(ref_audio, ref_sr),
                ref_text="",
                x_vector_only_mode=True
            )
        elif self.ref_hash:
            ref_hash = self.ref_hash
            prompt = self.prompt_items
        else:
            raise ValueError("Reference audio required for Base model")
        
        cached = self.audio_cache.get(text, ref_hash)
        if cached:
            return cached
        
        audio_chunks, sample_rate = self.model.generate_voice_clone(
            text=text,
            voice_clone_prompt=prompt,
            x_vector_only_mode=True,
        )
        
        wav_data = self._audio_to_wav(audio_chunks, sample_rate)
        self.audio_cache.put(text, ref_hash, wav_data)
        return wav_data
    
    def _audio_to_wav(self, audio_data, sample_rate: int = None) -> bytes:
        if sample_rate is None:
            sample_rate = self.sample_rate
        
        if isinstance(audio_data, list):
            audio_array = np.concatenate(audio_data) if len(audio_data) > 0 else np.array([])
        else:
            audio_array = audio_data
        
        if audio_array.dtype != np.int16:
            audio_array = (audio_array * 32767).astype(np.int16)
        
        wav_io = io.BytesIO()
        with wave.open(wav_io, 'wb') as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(sample_rate)
            wav_file.writeframes(audio_array.tobytes())
        
        return wav_io.getvalue()
    
    def get_model_info(self) -> dict:
        info = {
            'model_id': self.model_id,
            'sample_rate': self.sample_rate,
            'has_ref_audio': self.ref_audio is not None,
            'cache_stats': self.audio_cache.get_stats()
        }
        return info


class TTSServer:
    def __init__(
        self,
        model_id: str = "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
        port: int = 8766,
        host: str = "localhost",
        ref_audio_path: str = None,
        max_workers: int = 4
    ):
        self.port = port
        self.host = host
        self.ref_audio_path = ref_audio_path
        self.clients: Set[WebSocketServerProtocol] = set()
        self.is_running = True
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        
        print(f"Initializing TTS engine...")
        self.engine = TTSEngine(model_id, ref_audio_path)
        print(f"TTS Engine ready (max_workers={max_workers})")
    
    async def register(self, websocket: WebSocketServerProtocol):
        self.clients.add(websocket)
        print(f"Client connected: {websocket.remote_address}")
        
        await self.send_message(websocket, {
            "type": "connected",
            "model": self.engine.get_model_info()
        })
    
    async def unregister(self, websocket: WebSocketServerProtocol):
        self.clients.discard(websocket)
        print(f"Client disconnected: {websocket.remote_address}")
    
    async def send_message(self, websocket: WebSocketServerProtocol, message: dict):
        if websocket in self.clients:
            try:
                await websocket.send(json.dumps(message, ensure_ascii=False))
            except websockets.exceptions.ConnectionClosed:
                await self.unregister(websocket)
    
    async def handle_synthesize(self, websocket: WebSocketServerProtocol, 
                               text: str, ref_audio: str = None, chunk_index: int = -1):
        loop = asyncio.get_event_loop()
        try:
            print(f"Starting synthesis for chunk {chunk_index}: {text[:30]}...")
            
            audio_data = await loop.run_in_executor(
                self.executor,
                self.engine.synthesize,
                text,
                ref_audio
            )
            
            audio_base64 = base64.b64encode(audio_data).decode('utf-8')
            await self.send_message(websocket, {
                "type": "audio",
                "data": audio_base64,
                "format": "wav",
                "sample_rate": self.engine.sample_rate,
                "chunk_index": chunk_index,
                "is_first": True,
                "is_last": True
            })
            print(f"Sent audio for chunk {chunk_index} ({len(audio_data)} bytes)")
            
        except Exception as e:
            print(f"Synthesis error for chunk {chunk_index}: {e}")
            await self.send_message(websocket, {
                "type": "error",
                "message": str(e),
                "chunk_index": chunk_index
            })
    
    async def handle_message(self, websocket: WebSocketServerProtocol, message: dict):
        msg_type = message.get("type")
        
        if msg_type == "synthesize":
            text = message.get("text", "")
            ref_audio = message.get("ref_audio")
            chunk_index = message.get("chunk_index", -1)
            
            if not text:
                await self.send_message(websocket, {
                    "type": "error",
                    "message": "No text provided"
                })
                return
            
            asyncio.create_task(self.handle_synthesize(websocket, text, ref_audio, chunk_index))
        
        elif msg_type == "voices":
            await self.send_message(websocket, {
                "type": "voices",
                "voices": ["default"]
            })
        
        elif msg_type == "cache_stats":
            stats = self.engine.audio_cache.get_stats()
            await self.send_message(websocket, {
                "type": "cache_stats",
                **stats
            })
    
    async def handler(self, websocket: WebSocketServerProtocol):
        await self.register(websocket)
        try:
            async for raw_message in websocket:
                try:
                    if isinstance(raw_message, str):
                        message = json.loads(raw_message)
                        await self.handle_message(websocket, message)
                except json.JSONDecodeError:
                    print(f"Invalid JSON from {websocket.remote_address}")
                except Exception as e:
                    print(f"Error handling message: {e}")
        except websockets.exceptions.ConnectionClosed:
            pass
        finally:
            await self.unregister(websocket)
    
    async def start(self):
        print(f"Starting TTS server on {self.host}:{self.port}")
        async with websockets.serve(self.handler, self.host, self.port):
            await asyncio.Future()
    
    def run(self):
        asyncio.run(self.start())


def main():
    parser = argparse.ArgumentParser(description="TTS WebSocket Server (Optimized with Cache)")
    parser.add_argument("--model", "-m", default="Qwen/Qwen3-TTS-12Hz-0.6B-Base")
    parser.add_argument("--port", "-p", type=int, default=8766)
    parser.add_argument("--host", default="localhost")
    parser.add_argument("--ref-audio", "-r", default=None)
    parser.add_argument("--workers", "-w", type=int, default=4)
    
    args = parser.parse_args()
    
    server = TTSServer(
        model_id=args.model,
        port=args.port,
        host=args.host,
        ref_audio_path=args.ref_audio,
        max_workers=args.workers
    )
    
    def signal_handler(sig, frame):
        print("\nShutting down server...")
        server.is_running = False
        sys.exit(0)
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    server.run()


if __name__ == "__main__":
    main()