app.py

import os
import io
import asyncio
import time
import shutil
import numpy as np
import psutil
import soundfile as sf
import subprocess
import tempfile
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Generator
from contextlib import asynccontextmanager
import logging
import aiofiles 
import torch
from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Query
from fastapi.responses import Response, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import re
import hashlib
from functools import lru_cache
# Ensure the cloned neutts-air repository is in the path
import sys
sys.path.append(os.path.join(os.getcwd(), 'neutts-air'))
from neuttsair.neutts import NeuTTSAir

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("NeuTTS-API")

# --- Configuration & Utility Functions ---

# Explicitly use CPU as per Dockerfile and Hugging Face free tier compatibility
DEVICE = "cpu" 
# Configure Max Workers for concurrent synthesis threads (1-2 is safe for CPU-only)
MAX_WORKERS = 2
tts_executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)
SAMPLE_RATE = 24000
CLEANUP_THRESHOLD = 300 # 1 hour in seconds
TEMP_AUDIO_DIR = "temp_audio"
GENERATED_AUDIO_DIR = "generated_audio"
os.makedirs(TEMP_AUDIO_DIR, exist_ok=True)
os.makedirs(GENERATED_AUDIO_DIR, exist_ok=True)

class TTSRequestModel(BaseModel):
    """Model for non-file inputs to synthesis and streaming."""
    text: str = Field(..., min_length=1, max_length=1000)
    speed: float = Field(default=1.0, ge=0.5, le=2.0)
    output_format: str = Field(default="wav", pattern="^(wav|mp3|flac)$")


async def convert_to_wav_in_memory(upload_file: UploadFile) -> io.BytesIO:
    """
    Converts uploaded audio to a 24kHz WAV in memory using FFmpeg pipes.
    This avoids all intermediate disk I/O for maximum speed.
    """
    ffmpeg_command = [
        "ffmpeg",
        "-i", "pipe:0",      # Read from stdin
        "-f", "wav",
        "-ar", str(SAMPLE_RATE),
        "-ac", "1",
        "-c:a", "pcm_s16le",
        "pipe:1"             # Write to stdout
    ]
    
    # Start the subprocess with pipes for stdin, stdout, and stderr
    proc = await asyncio.create_subprocess_exec(
        *ffmpeg_command,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )

    # Stream the uploaded file data into ffmpeg's stdin
    # and capture the resulting WAV data from its stdout
    wav_data, stderr_data = await proc.communicate(input=await upload_file.read())

    if proc.returncode != 0:
        error_message = stderr_data.decode()
        logger.error(f"In-memory conversion failed: {error_message}")
        # Provide the last line of the FFmpeg error to the user
        error_detail = error_message.splitlines()[-1] if error_message else "Unknown FFmpeg error."
        raise HTTPException(status_code=400, detail=f"Audio format conversion failed: {error_detail}")
    
    logger.info("In-memory FFmpeg conversion successful.")
    # Return the raw WAV data in a BytesIO buffer, ready for the model
    return io.BytesIO(wav_data)
# --- Model Wrapper and Logic ---

class NeuTTSWrapper:
    def __init__(self, device: str = "cpu"):
        self.tts_model = None
        self.device = device
        self.load_model()

    def load_model(self):
        try:
            logger.info(f"Loading NeuTTSAir model on device: {self.device}")
            # Ensure we respect the CPU configuration
            self.tts_model = NeuTTSAir(backbone_device=self.device, codec_device=self.device)
            logger.info("✅ NeuTTSAir model loaded successfully.")
        except Exception as e:
            logger.error(f"❌ Model loading failed: {e}")
            raise

    def _convert_to_streamable_format(self, audio_data: np.ndarray, audio_format: str) -> bytes:
        """Converts NumPy audio array to streamable bytes in the specified format."""
        audio_buffer = io.BytesIO()
        try:
            sf.write(audio_buffer, audio_data, SAMPLE_RATE, format=audio_format)
        except Exception as e:
            logger.error(f"Failed to write audio data to format {audio_format}: {e}")
            raise
        audio_buffer.seek(0)
        return audio_buffer.read()

    def _split_text_into_chunks(self, text: str) -> list[str]:
        """
        Splits text into sentences OR clauses using a robust regex. 
        This is fast, library-free, and now handles commas.
        """
        # This regex now finds all sequences of characters that are not a sentence-ending
        # or clause-ending punctuation mark, followed by that punctuation.
        # The only change is adding ',' to the character sets.
        chunks = re.findall(r'[^.,!?]+[.,!?]*', text)
        return [c.strip() for c in chunks if c.strip()]

    @lru_cache(maxsize=32)
    def _get_or_create_reference_encoding(self, audio_content_hash: str, audio_bytes: bytes) -> torch.Tensor:
        """
        Caches the expensive reference encoding operation using an in-memory LRU cache.
        The hash of the audio content is the key.
        """
        logger.info(f"Cache miss for hash: {audio_content_hash[:10]}... Encoding new reference.")
        # The model's encode_reference can take a file-like object (BytesIO)
        return self.tts_model.encode_reference(io.BytesIO(audio_bytes))

    def generate_speech_blocking(self, text: str, ref_audio_bytes: bytes, reference_text: str) -> np.ndarray:
        """Blocking synthesis using cached reference encoding."""
        # 1. Hash the audio bytes to get a cache key
        audio_hash = hashlib.sha256(ref_audio_bytes).hexdigest()
        
        # 2. Get the encoding from the cache (or create it if new)
        ref_s = self._get_or_create_reference_encoding(audio_hash, ref_audio_bytes)
        
        # 3. Infer full text
        with torch.no_grad():
            audio = self.tts_model.infer(text, ref_s, reference_text)
        return audio

    def stream_speech_blocking(self, text: str, ref_audio_bytes: bytes, reference_text: str, speed: float, audio_format: str) -> Generator[bytes, None, None]:
        """Sentence-by-Sentence Streaming using cached reference encoding."""
        logger.info(f"Starting streaming synthesis for text length: {len(text)}")
        
        # 1. Hash the audio bytes once
        audio_hash = hashlib.sha256(ref_audio_bytes).hexdigest()

        # 2. Get the reference encoding from cache, once for the whole stream
        ref_s = self._get_or_create_reference_encoding(audio_hash, ref_audio_bytes)
        
        # 3. Split text using the new regex method
        sentences = self._split_text_into_chunks(text)
        
        # 4. Stream chunks
        for i, sentence in enumerate(sentences):
            if not sentence.strip():
                continue
            
            logger.debug(f"Generating streaming chunk {i+1}: '{sentence[:30]}...'")
            
            with torch.no_grad():
                audio_chunk = self.tts_model.infer(sentence, ref_s, reference_text)
            
            yield self._convert_to_streamable_format(audio_chunk, audio_format)
            
        logger.info("Streaming synthesis complete.")

# --- Asynchronous Offloading ---

async def run_blocking_task_async(func, *args, **kwargs):
    """Offloads a blocking function call to the ThreadPoolExecutor."""
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(
        tts_executor,
        lambda: func(*args, **kwargs)
    )

async def save_upload_file_async(upload_file: UploadFile) -> str:
    """Asynchronously saves the UploadFile to disk."""
    temp_filename = os.path.join(TEMP_AUDIO_DIR, f"{time.time()}_{upload_file.filename}")
    try:
        # Use asyncio to read the file chunks in a non-blocking manner
        async with aiofiles.open(temp_filename, 'wb') as out_file:
            while content := await upload_file.read(1024 * 1024):
                await out_file.write(content)
        return temp_filename
    except Exception as e:
        logger.error(f"Error saving file: {e}")
        raise HTTPException(status_code=500, detail="Could not save reference audio file")

# --- FastAPI Lifespan Manager (Kokoro Feature) ---

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Modern lifespan management: initialize model on startup, shutdown executor."""
    try:
        app.state.tts_wrapper = NeuTTSWrapper(device=DEVICE)
    except Exception as e:
        logger.error(f"Fatal startup error: {e}")
        # Terminate the application if the model can't load
        tts_executor.shutdown(wait=False) 
        raise RuntimeError("Model initialization failed.")
    
    yield # Application serves requests
    
    # Shutdown
    logger.info("Shutting down ThreadPoolExecutor.")
    tts_executor.shutdown(wait=False)

# --- FastAPI Application Setup ---
app = FastAPI(
    title="NeuTTS Air Instant Cloning API", 
    version="2.0.0-PROD-ENHANCED", 
    docs_url="/docs", 
    lifespan=lifespan
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- New Endpoints and Enhancements ---

@app.get("/")
async def root():
    return {"message": "NeuTTS Air API v2.0 - Ready for Instant Voice Cloning"}

@app.get("/health")
async def health_check():
    """Enhanced health check (Kokoro Feature + Original Metrics)"""
    mem = psutil.virtual_memory()
    disk = psutil.disk_usage('/')
    
    return {
        "status": "healthy",
        "model_loaded": hasattr(app.state, 'tts_wrapper') and app.state.tts_wrapper.tts_model is not None,
        "device": DEVICE,
        "concurrency_limit": MAX_WORKERS,
        "memory_usage": {
            "total_gb": round(mem.total / (1024**3), 2),
            "used_percent": mem.percent
        },
        "disk_usage": {
            "total_gb": round(disk.total / (1024**3), 2),
            "used_percent": disk.percent
        }
    }

@app.delete("/cleanup")
async def cleanup_files():
    """Maintenance endpoint to remove old generated and temporary files."""
    await run_blocking_task_async(cleanup_files_blocking)
    return {"message": "Cleanup initiated successfully."}

def cleanup_files_blocking():
    """Blocking file cleanup logic (original NeuTTS feature)."""
    now = time.time()
    deleted_count = 0
    
    for directory in [GENERATED_AUDIO_DIR, TEMP_AUDIO_DIR]:
        for filename in os.listdir(directory):
            filepath = os.path.join(directory, filename)
            if os.path.isfile(filepath):
                try:
                    # Original cleanup logic: delete if older than CLEANUP_THRESHOLD
                    if now - os.path.getctime(filepath) > CLEANUP_THRESHOLD:
                        os.remove(filepath)
                        deleted_count += 1
                except Exception as e:
                    logger.warning(f"Failed to delete {filepath}: {e}")
    
    logger.info(f"Cleanup completed: {deleted_count} files removed.")
    return deleted_count


# --- Core Synthesis Endpoints ---

@app.post("/synthesize", response_class=Response)
async def text_to_speech(
    text: str = Form(...),
    reference_text: str = Form(...),
    speed: float = Form(1.0, ge=0.5, le=2.0),
    output_format: str = Form("wav", pattern="^(wav|mp3|flac)$"),
    reference_audio: UploadFile = File(...)):
    """
    MAXIMUM SPEED TTS endpoint with full parallel processing.
    """
    if not hasattr(app.state, 'tts_wrapper'):
        raise HTTPException(status_code=503, detail="Service unavailable: Model not loaded")
    
    start_time = time.time()
    try:
        # ✅ PARALLEL STEP 1: Convert audio AND split text concurrently
        converted_wav_buffer, sentences = await asyncio.gather(
            convert_to_wav_in_memory(reference_audio),
            asyncio.get_event_loop().run_in_executor(
                tts_executor, 
                app.state.tts_wrapper._split_text_into_chunks, 
                text
            )
        )
        
        ref_audio_bytes = converted_wav_buffer.getvalue()
        audio_hash = hashlib.sha256(ref_audio_bytes).hexdigest()
        logger.info(f"🚀 MAX PARALLEL: Processing {len(sentences)} chunks")

        # ✅ PARALLEL STEP 2: Get reference encoding
        ref_s = await asyncio.get_event_loop().run_in_executor(
            tts_executor, 
            app.state.tts_wrapper._get_or_create_reference_encoding,
            audio_hash,
            ref_audio_bytes
        )
        
        # ✅ MAX PARALLEL STEP 3: Process ALL chunks simultaneously
        loop = asyncio.get_event_loop()
        
        def process_single_chunk(sentence_text):
            with torch.no_grad():
                return app.state.tts_wrapper.tts_model.infer(sentence_text, ref_s, reference_text)
        
        # Schedule ALL chunks in parallel (limited by MAX_WORKERS)
        tasks = []
        for sentence in sentences:
            task = loop.run_in_executor(tts_executor, process_single_chunk, sentence)
            tasks.append(task)
        
        # Wait for ALL chunks to complete
        chunk_audios = await asyncio.gather(*tasks)
        
        # ✅ Combine all audio chunks (fast numpy concatenation)
        combined_audio = np.concatenate(chunk_audios) if chunk_audios else np.array([])
        
        # ✅ PARALLEL STEP 4: Convert format while calculating stats
        audio_bytes, audio_duration = await asyncio.gather(
            asyncio.get_event_loop().run_in_executor(
                tts_executor,
                app.state.tts_wrapper._convert_to_streamable_format,
                combined_audio,
                output_format
            ),
            asyncio.get_event_loop().run_in_executor(
                tts_executor,
                lambda: len(combined_audio) / SAMPLE_RATE
            )
        )
                
        processing_time = time.time() - start_time
        
        logger.info(f"✅ MAX SPEED Synthesis: {processing_time:.2f}s for {audio_duration:.2f}s audio ({len(sentences)} chunks)")
        
        return Response(
            content=audio_bytes,
            media_type=f"audio/{'mpeg' if output_format == 'mp3' else output_format}",
            headers={
                "Content-Disposition": f"attachment; filename=tts_output.{output_format}",
                "X-Processing-Time": f"{processing_time:.2f}s",
                "X-Audio-Duration": f"{audio_duration:.2f}s",
                "X-Parallel-Chunks": str(len(sentences)),
                "X-Speed-Ratio": f"{audio_duration/processing_time:.2f}x"  # Real-time factor
            }
        )
    except Exception as e:
        logger.error(f"Synthesis error: {e}")
        if isinstance(e, HTTPException):
             raise
        raise HTTPException(status_code=500, detail=f"Synthesis failed: {e}")
@app.post("/synthesize/stream")
async def stream_text_to_speech_cloning(
    text: str = Form(..., min_length=1, max_length=5000),
    reference_text: str = Form(...),
    speed: float = Form(1.0, ge=0.5, le=2.0),
    output_format: str = Form("mp3", pattern="^(wav|mp3|flac)$"),
    reference_audio: UploadFile = File(...)):
    """
    TRUE Real-Time Streaming with 2 workers: Optimized for continuous audio.
    """
    if not hasattr(app.state, 'tts_wrapper'):
        raise HTTPException(status_code=503, detail="Service unavailable: Model not loaded")

    async def stream_generator():
        loop = asyncio.get_event_loop()
        
        # ✅ Perfect queue size for 2 workers
        q = asyncio.Queue(maxsize=3)  # Store 3 ready chunks for smooth streaming

        async def producer():
            try:
                converted_wav_buffer = await convert_to_wav_in_memory(reference_audio)
                ref_audio_bytes = converted_wav_buffer.getvalue()
                audio_hash = hashlib.sha256(ref_audio_bytes).hexdigest()
                
                # Get reference encoding (uses 1 worker temporarily)
                ref_s = await loop.run_in_executor(
                    tts_executor, 
                    app.state.tts_wrapper._get_or_create_reference_encoding, 
                    audio_hash, 
                    ref_audio_bytes
                )
                
                sentences = app.state.tts_wrapper._split_text_into_chunks(text)
                logger.info(f"Streaming {len(sentences)} chunks with 2 workers")
                
                def process_chunk(sentence_text):
                    with torch.no_grad():
                        audio_chunk = app.state.tts_wrapper.tts_model.infer(sentence_text, ref_s, reference_text)
                    return app.state.tts_wrapper._convert_to_streamable_format(audio_chunk, output_format)

                # ✅ SCHEDULE ALL TASKS IMMEDIATELY
                for sentence in sentences:
                    task = loop.run_in_executor(tts_executor, process_chunk, sentence)
                    await q.put(task)  # Queue futures immediately
                    
            except Exception as e:
                logger.error(f"Error in producer task: {e}")
                await q.put(e)
            finally:
                await q.put(None)  # Signal end of tasks

        producer_task = asyncio.create_task(producer())

        # ✅ EFFICIENT CONSUMER for 2 workers
        pending_tasks = set()
        completed_count = 0
        total_chunks = len(app.state.tts_wrapper._split_text_into_chunks(text))
        
        while completed_count < total_chunks:
            # Get next item from queue
            result = await q.get()
            
            if isinstance(result, Exception):
                logger.error(f"Terminating stream due to producer error: {result}")
                raise result
            
            if result is None:
                break  # No more tasks coming
                
            # ✅ Process this chunk immediately
            chunk_bytes = await result
            yield chunk_bytes
            completed_count += 1
            
            # ✅ Check if we can process next chunk without waiting
            # This ensures continuous streaming
            if completed_count < total_chunks and not q.empty():
                continue  # Immediately process next ready chunk
                
        await producer_task

    return StreamingResponse(
        stream_generator(),
        media_type=f"audio/{'mpeg' if output_format == 'mp3' else output_format}"
    )

@app.get("/audio/{filename}")
async def get_audio(filename: str):
    """Original NeuTTS feature to serve generated audio files."""
    file_path = os.path.join(GENERATED_AUDIO_DIR, filename)
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="Audio file not found")
    
    return Response(
        content=open(file_path, "rb").read(),
        media_type=f"audio/{filename.split('.')[-1]}", # Simple media type detection
        headers={"Content-Disposition": f"attachment; filename={filename}"}
    )