server.py

"""
Studio Server - AI-powered studio utilities for video production.

Provides REST API endpoints for:
- TTS: Text-to-speech with voice cloning (Qwen3-TTS)
- Face: Face embedding extraction for IP-Adapter (InsightFace)
- Transcription: Audio transcription with word timings (Whisper)

All backends are modular and can be swapped via environment variables.
"""

import io
import os
from typing import Optional
from contextlib import asynccontextmanager

import soundfile as sf
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel

from backends.tts import TTSBackend, get_tts_backend, TTS_BACKENDS
from backends.face import FaceBackend, get_face_backend, FACE_BACKENDS
from backends.transcription import TranscriptionBackend, get_transcription_backend, TRANSCRIPTION_BACKENDS


# =============================================================================
# Backend Instances
# =============================================================================

tts_backend: Optional[TTSBackend] = None
face_backend: Optional[FaceBackend] = None
transcription_backend: Optional[TranscriptionBackend] = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load backends on startup based on environment config."""
    global tts_backend, face_backend, transcription_backend

    # Load TTS backend (always enabled)
    tts_name = os.environ.get("TTS_BACKEND", "qwen3-tts")
    print(f"\n=== Loading TTS backend: {tts_name} ===")
    tts_backend = get_tts_backend(tts_name)
    tts_backend.load()

    # Load Face backend (optional, enabled by default)
    if os.environ.get("FACE_ENABLED", "true").lower() == "true":
        face_name = os.environ.get("FACE_BACKEND", "insightface")
        print(f"\n=== Loading Face backend: {face_name} ===")
        face_backend = get_face_backend(face_name)
        face_backend.load()

    # Load Transcription backend (optional, enabled by default)
    if os.environ.get("TRANSCRIPTION_ENABLED", "true").lower() == "true":
        transcription_name = os.environ.get("TRANSCRIPTION_BACKEND", "whisper")
        print(f"\n=== Loading Transcription backend: {transcription_name} ===")
        transcription_backend = get_transcription_backend(transcription_name)
        transcription_backend.load()

    print("\n=== Studio Server Ready ===\n")
    yield

    # Cleanup
    if tts_backend:
        tts_backend.unload()
    if face_backend:
        face_backend.unload()
    if transcription_backend:
        transcription_backend.unload()


app = FastAPI(
    title="Studio Server",
    description="AI-powered studio utilities for video production: TTS, face embedding, transcription",
    version="0.4.0",
    lifespan=lifespan,
)


# =============================================================================
# Health & Info Endpoints
# =============================================================================

@app.get("/health")
async def health():
    """Health check endpoint."""
    backends = {}
    if tts_backend:
        backends["tts"] = tts_backend.get_info()
    if face_backend:
        backends["face"] = face_backend.get_info()
    if transcription_backend:
        backends["transcription"] = transcription_backend.get_info()

    return {
        "status": "ok",
        "backends": backends,
    }


@app.get("/v1/models")
async def list_models():
    """List available backends by category."""
    return {
        "tts": list(TTS_BACKENDS.keys()),
        "face": list(FACE_BACKENDS.keys()),
        "transcription": list(TRANSCRIPTION_BACKENDS.keys()),
    }


# =============================================================================
# TTS Endpoints
# =============================================================================

@app.get("/v1/tts/speakers")
async def list_speakers():
    """List available preset speakers for basic TTS."""
    if tts_backend is None:
        raise HTTPException(status_code=503, detail="TTS backend not loaded")
    return {"speakers": tts_backend.get_speakers()}


@app.post("/v1/tts/extract")
async def extract_voice_prompt(
    ref_audio: UploadFile = File(..., description="Reference audio for voice extraction"),
    ref_text: Optional[str] = Form(None, description="Transcript of reference audio"),
):
    """
    Extract a reusable voice prompt from reference audio.

    The returned `voice_prompt` can be cached and reused with `/v1/tts/synthesize`
    to avoid re-processing the reference audio on every request.
    """
    if tts_backend is None:
        raise HTTPException(status_code=503, detail="TTS backend not loaded")

    try:
        audio_bytes = await ref_audio.read()
        audio_buffer = io.BytesIO(audio_bytes)
        audio_data, sample_rate = sf.read(audio_buffer)
        ref_audio_data = (audio_data, sample_rate)

        voice_prompt = tts_backend.extract_voice_prompt(
            ref_audio=ref_audio_data,
            ref_text=ref_text,
        )

        return JSONResponse({
            "voice_prompt": voice_prompt,
            "format": "base64-torch",
            "ref_text": ref_text,
        })

    except NotImplementedError as e:
        raise HTTPException(status_code=501, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/v1/tts/synthesize")
async def synthesize_speech(
    text: str = Form(..., description="Text to synthesize"),
    language: str = Form("English", description="Target language"),
    speaker: Optional[str] = Form(None, description="Preset speaker for basic TTS"),
    ref_audio: Optional[UploadFile] = File(None, description="Reference audio for voice cloning"),
    ref_text: Optional[str] = Form(None, description="Transcript of reference audio"),
    voice_prompt: Optional[str] = Form(None, description="Pre-extracted voice prompt"),
    speed: float = Form(1.0, description="Speech speed multiplier"),
):
    """
    Synthesize speech from text.

    **Basic TTS:** Just provide `text` and optionally `speaker`.

    **Voice Cloning (on-the-fly):** Provide `ref_audio` and `ref_text`.

    **Voice Cloning (cached):** Provide `voice_prompt` from `/v1/tts/extract`.
    """
    if tts_backend is None:
        raise HTTPException(status_code=503, detail="TTS backend not loaded")

    try:
        if voice_prompt:
            wav_bytes, sample_rate = tts_backend.synthesize_with_prompt(
                text=text,
                voice_prompt=voice_prompt,
                language=language,
                speed=speed,
            )
        elif ref_audio:
            audio_bytes = await ref_audio.read()
            audio_buffer = io.BytesIO(audio_bytes)
            audio_data, sample_rate = sf.read(audio_buffer)
            ref_audio_data = (audio_data, sample_rate)

            wav_bytes, sample_rate = tts_backend.synthesize(
                text=text,
                language=language,
                speaker=speaker,
                ref_audio=ref_audio_data,
                ref_text=ref_text,
                speed=speed,
            )
        else:
            wav_bytes, sample_rate = tts_backend.synthesize(
                text=text,
                language=language,
                speaker=speaker,
                ref_audio=None,
                ref_text=None,
                speed=speed,
            )

        return StreamingResponse(
            io.BytesIO(wav_bytes),
            media_type="audio/wav",
            headers={
                "Content-Disposition": "attachment; filename=speech.wav",
                "X-Sample-Rate": str(sample_rate),
            }
        )

    except NotImplementedError as e:
        raise HTTPException(status_code=501, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# =============================================================================
# Face Embedding Endpoints
# =============================================================================

@app.post("/v1/face/embed")
async def extract_face_embedding(
    image: UploadFile = File(..., description="Image containing a face"),
    return_bbox: bool = Form(False, description="Include bounding box in response"),
):
    """
    Extract face embedding from an image.

    Returns a 512-dimensional embedding compatible with IP-Adapter FaceID
    for maintaining face consistency across video generation (PerformerDNA).
    """
    if face_backend is None:
        raise HTTPException(status_code=503, detail="Face backend not loaded")

    try:
        image_bytes = await image.read()
        result = face_backend.extract_embedding(image_bytes, return_bbox=return_bbox)
        return JSONResponse(result)

    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/v1/face/embed-all")
async def extract_all_face_embeddings(
    image: UploadFile = File(..., description="Image containing faces"),
    max_faces: int = Form(5, description="Maximum number of faces to extract"),
):
    """
    Extract face embeddings for all faces in an image.

    Useful when processing reference images that may contain multiple people.
    Returns faces sorted by size (largest first).
    """
    if face_backend is None:
        raise HTTPException(status_code=503, detail="Face backend not loaded")

    try:
        image_bytes = await image.read()
        results = face_backend.extract_embeddings(image_bytes, max_faces=max_faces)
        return JSONResponse({"faces": results, "count": len(results)})

    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


class CompareEmbeddingsRequest(BaseModel):
    """Request body for comparing face embeddings."""
    embedding1: str
    embedding2: str


@app.post("/v1/face/compare")
async def compare_face_embeddings(request: CompareEmbeddingsRequest):
    """
    Compare two face embeddings for similarity.

    Returns a cosine similarity score between 0 and 1.
    Scores above 0.6 typically indicate the same person.
    """
    if face_backend is None:
        raise HTTPException(status_code=503, detail="Face backend not loaded")

    try:
        similarity = face_backend.compare_embeddings(
            request.embedding1,
            request.embedding2,
        )
        return JSONResponse({
            "similarity": similarity,
            "same_person": similarity > 0.6,
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# =============================================================================
# Transcription Endpoints
# =============================================================================

@app.post("/v1/transcribe")
async def transcribe_audio(
    audio: UploadFile = File(..., description="Audio file to transcribe"),
    language: Optional[str] = Form(None, description="Language code (auto-detect if empty)"),
    word_timestamps: bool = Form(True, description="Include word-level timestamps"),
):
    """
    Transcribe audio to text with word-level timestamps.

    Returns the full transcript and per-word timing information
    for lip-sync alignment.
    """
    if transcription_backend is None:
        raise HTTPException(status_code=503, detail="Transcription backend not loaded")

    try:
        audio_bytes = await audio.read()
        result = transcription_backend.transcribe(
            audio_bytes,
            language=language,
            word_timestamps=word_timestamps,
        )
        return JSONResponse(result.to_dict())

    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# =============================================================================
# Main
# =============================================================================

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)