app.py

from __future__ import annotations

import asyncio
import os
import tempfile
import time
import unicodedata
from pathlib import Path
from typing import Annotated

import httpx
from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, Response
from fastapi.staticfiles import StaticFiles


ROOT = Path(__file__).resolve().parent
BACKEND_URL = os.getenv("BACKEND_URL", "http://127.0.0.1:8010")
MAX_UPLOAD_BYTES = 15 * 1024 * 1024
ALLOWED_AUDIO_SUFFIXES = {".wav", ".mp3", ".flac", ".m4a", ".ogg"}
MAX_TEXT_UNITS = int(os.getenv("AUDIO8_TTS_MAX_TEXT_UNITS", "150"))
MAX_RAW_TEXT_CHARS = int(os.getenv("AUDIO8_TTS_MAX_RAW_TEXT_CHARS", "1000"))
MAX_NEW_TOKENS = int(os.getenv("AUDIO8_TTS_MAX_NEW_TOKENS", "1024"))
GENERATION_LIMIT = asyncio.Semaphore(int(os.getenv("UI_MAX_CONCURRENCY", "1")))

ENGLISH_REFERENCE_TEXT = (
    "hello nice to meet you, what would you like to talk about todat"
)
CHINESE_REFERENCE_TEXT = "你好，我是小周，很高兴认识你"

EXAMPLES = {
    "clara": {
        "id": "clara",
        "name": "Clara",
        "locale": "English",
        "tone": "Female",
        "file": "en_female_clara.wav",
        "transcript": ENGLISH_REFERENCE_TEXT,
    },
    "iris": {
        "id": "iris",
        "name": "Iris",
        "locale": "English",
        "tone": "Female",
        "file": "en_female_iris.wav",
        "transcript": ENGLISH_REFERENCE_TEXT,
    },
    "arthur": {
        "id": "arthur",
        "name": "Arthur",
        "locale": "English",
        "tone": "Male",
        "file": "en_male_arthur.wav",
        "transcript": ENGLISH_REFERENCE_TEXT,
    },
    "mia": {
        "id": "mia",
        "name": "Mia",
        "locale": "中文",
        "tone": "女声",
        "file": "zh_female_mia.wav",
        "transcript": CHINESE_REFERENCE_TEXT,
    },
    "ben": {
        "id": "ben",
        "name": "Ben",
        "locale": "中文",
        "tone": "男声",
        "file": "zh_male_ben.wav",
        "transcript": CHINESE_REFERENCE_TEXT,
    },
    "sophie": {
        "id": "sophie",
        "name": "Sophie",
        "locale": "中英双语",
        "tone": "女声",
        "file": "zh_en_female_sophie.wav",
        "transcript": CHINESE_REFERENCE_TEXT,
    },
}

app = FastAPI(title="Audio8 TTS Preview 0.6B", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)
app.mount("/static", StaticFiles(directory=ROOT / "static"), name="static")


@app.get("/assets/{asset_name}", include_in_schema=False)
async def asset(asset_name: str) -> FileResponse:
    if asset_name != "audio8-logo.jpeg":
        raise HTTPException(status_code=404, detail="Asset not found")
    return FileResponse(
        ROOT / "assets" / asset_name,
        media_type="image/jpeg",
        headers={"Cache-Control": "public, max-age=86400"},
    )


async def _backend_status() -> tuple[bool, dict]:
    try:
        async with httpx.AsyncClient(timeout=3.0) as client:
            response = await client.get(f"{BACKEND_URL}/health")
        payload = response.json()
        return response.status_code == 200, payload
    except (httpx.HTTPError, ValueError):
        return False, {}


def _normalize_speech_text(text: str) -> str:
    separator = "\ue000"
    text = text.replace("\\n", separator).replace("\\r", separator).replace("\\t", separator)

    def is_latin_or_number(character: str) -> bool:
        name = unicodedata.name(character, "")
        return character.isnumeric() or (character.isalpha() and "LATIN" in name)

    def next_visible_character(start: int) -> str:
        for candidate in text[start:]:
            if candidate == separator or candidate.isspace():
                continue
            if unicodedata.category(candidate) in {"Cc", "Cf", "Cs", "Co", "Cn"}:
                continue
            return candidate
        return ""

    cleaned: list[str] = []
    for index, character in enumerate(text):
        category = unicodedata.category(character)
        if character == separator or (character.isspace() and category == "Cc"):
            previous = next((item for item in reversed(cleaned) if not item.isspace()), "")
            following = next_visible_character(index + 1)
            if is_latin_or_number(previous) and is_latin_or_number(following):
                cleaned.append(" ")
            continue
        if category in {"Cc", "Cf", "Cs", "Co", "Cn"}:
            continue
        cleaned.append(" " if character.isspace() else character)
    return " ".join("".join(cleaned).split())


def _is_cjk_character(character: str) -> bool:
    codepoint = ord(character)
    return (
        0x3400 <= codepoint <= 0x4DBF
        or 0x4E00 <= codepoint <= 0x9FFF
        or 0xF900 <= codepoint <= 0xFAFF
        or 0x3040 <= codepoint <= 0x30FF
        or 0xAC00 <= codepoint <= 0xD7AF
    )


def _count_speech_units(text: str) -> int:
    units = 0
    in_latin_word = False
    for character in text:
        if _is_cjk_character(character):
            units += 1
            in_latin_word = False
            continue

        name = unicodedata.name(character, "")
        if character.isnumeric() or (character.isalpha() and "LATIN" in name):
            if not in_latin_word:
                units += 1
            in_latin_word = True
        elif character in {"'", "\u2019", "-"} and in_latin_word:
            continue
        else:
            in_latin_word = False
            if character.isalpha() or character.isnumeric():
                units += 1
    return units


def _validate_speech_text(text: str) -> str:
    text = _normalize_speech_text(text)
    if len(text) > MAX_RAW_TEXT_CHARS:
        raise HTTPException(status_code=400, detail="Speech text is too long")
    if not text:
        raise HTTPException(status_code=400, detail="Text must not be empty")
    units = _count_speech_units(text)
    if units == 0:
        raise HTTPException(status_code=400, detail="Text must contain readable characters")
    if units > MAX_TEXT_UNITS:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Text must be {MAX_TEXT_UNITS} Chinese characters or "
                "English words or fewer"
            ),
        )
    return text


def _validate_speech_payload(payload: dict) -> dict:
    text = payload.get("input")
    if not isinstance(text, str):
        raise HTTPException(status_code=400, detail="Text must not be empty")
    text = _validate_speech_text(text)

    max_new_tokens = payload.get("max_new_tokens", MAX_NEW_TOKENS)
    if isinstance(max_new_tokens, bool):
        raise HTTPException(status_code=400, detail="Max tokens must be an integer")
    try:
        max_new_tokens = int(max_new_tokens)
    except (TypeError, ValueError) as exc:
        raise HTTPException(status_code=400, detail="Max tokens must be an integer") from exc
    if not 32 <= max_new_tokens <= MAX_NEW_TOKENS:
        raise HTTPException(
            status_code=400,
            detail=f"Max tokens must be between 32 and {MAX_NEW_TOKENS}",
        )

    return {**payload, "input": text, "max_new_tokens": max_new_tokens}


async def _generate(payload: dict) -> httpx.Response:
    payload = _validate_speech_payload(payload)
    async with GENERATION_LIMIT:
        async with httpx.AsyncClient(timeout=httpx.Timeout(600.0)) as client:
            return await client.post(f"{BACKEND_URL}/v1/audio/speech", json=payload)


async def _store_upload(upload: UploadFile) -> Path:
    suffix = Path(upload.filename or "reference.wav").suffix.lower()
    if suffix not in ALLOWED_AUDIO_SUFFIXES:
        raise HTTPException(status_code=400, detail="Unsupported reference audio format")

    descriptor, raw_path = tempfile.mkstemp(prefix="audio8-reference-", suffix=suffix)
    path = Path(raw_path)
    total = 0
    try:
        with os.fdopen(descriptor, "wb") as output:
            while chunk := await upload.read(1024 * 1024):
                total += len(chunk)
                if total > MAX_UPLOAD_BYTES:
                    raise HTTPException(
                        status_code=413,
                        detail="Reference audio must be 15 MB or smaller",
                    )
                output.write(chunk)
        return path
    except Exception:
        path.unlink(missing_ok=True)
        raise
    finally:
        await upload.close()


@app.get("/", include_in_schema=False)
async def index() -> FileResponse:
    return FileResponse(ROOT / "static" / "index.html")


@app.get("/api/examples")
async def list_examples() -> list[dict]:
    return [
        {
            **example,
            "audio_url": f"/examples/{example['id']}",
        }
        for example in EXAMPLES.values()
    ]


@app.get("/examples/{example_id}", include_in_schema=False)
async def example_audio(example_id: str) -> FileResponse:
    example = EXAMPLES.get(example_id)
    if example is None:
        raise HTTPException(status_code=404, detail="Voice example not found")
    return FileResponse(
        ROOT / "examples" / example["file"],
        media_type="audio/wav",
        headers={"Cache-Control": "public, max-age=86400"},
    )


@app.get("/api/status")
async def api_status() -> dict:
    ready, details = await _backend_status()
    return {
        "state": "ready" if ready else "warming",
        "model": "Audio8/Audio8-TTS-Preview-0.6b",
        "engine": "SGLang-Omni 0.1.0 / SGLang 0.5.8",
        "details": details,
    }


@app.post("/api/generate")
async def generate_speech(
    text: Annotated[str, Form()],
    reference_text: Annotated[str, Form()],
    example_id: Annotated[str | None, Form()] = None,
    reference_audio: Annotated[UploadFile | None, File()] = None,
    temperature: Annotated[float, Form()] = 0.8,
    top_p: Annotated[float, Form()] = 0.95,
    top_k: Annotated[int, Form()] = 50,
    max_new_tokens: Annotated[int, Form()] = 1024,
) -> Response:
    text = _validate_speech_text(text)
    reference_text = " ".join(reference_text.split())
    if not reference_text:
        raise HTTPException(status_code=400, detail="Reference transcript is required")
    if not 0 <= temperature <= 2:
        raise HTTPException(status_code=400, detail="Temperature must be between 0 and 2")
    if not 0 < top_p <= 1:
        raise HTTPException(status_code=400, detail="Top P must be between 0 and 1")
    if not 1 <= top_k <= 200:
        raise HTTPException(status_code=400, detail="Top K must be between 1 and 200")
    if not 32 <= max_new_tokens <= MAX_NEW_TOKENS:
        raise HTTPException(
            status_code=400,
            detail=f"Max tokens must be between 32 and {MAX_NEW_TOKENS}",
        )

    temporary_path: Path | None = None
    if reference_audio is not None and reference_audio.filename:
        temporary_path = await _store_upload(reference_audio)
        reference_path = temporary_path
    elif example_id and example_id in EXAMPLES:
        reference_path = ROOT / "examples" / EXAMPLES[example_id]["file"]
    else:
        raise HTTPException(status_code=400, detail="Select or upload a reference voice")

    payload = {
        "model": "audio8/tts-0.6b",
        "input": text,
        "response_format": "wav",
        "max_new_tokens": max_new_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "references": [
            {
                "audio_path": str(reference_path),
                "text": reference_text,
            }
        ],
    }

    started = time.perf_counter()
    try:
        response = await _generate(payload)
    except httpx.ConnectError as exc:
        raise HTTPException(status_code=503, detail="Model is still warming up") from exc
    except httpx.TimeoutException as exc:
        raise HTTPException(status_code=504, detail="Generation timed out") from exc
    finally:
        if temporary_path is not None:
            temporary_path.unlink(missing_ok=True)

    if response.status_code != 200:
        try:
            detail = response.json().get("detail", response.text)
        except ValueError:
            detail = response.text
        raise HTTPException(status_code=response.status_code, detail=detail)

    headers = {
        "Content-Disposition": 'attachment; filename="audio8-clone.wav"',
        "X-Generation-Duration-Ms": str(round((time.perf_counter() - started) * 1000)),
    }
    for name in ("x-prompt-tokens", "x-completion-tokens", "x-engine-time"):
        if name in response.headers:
            headers[name] = response.headers[name]
    return Response(content=response.content, media_type="audio/wav", headers=headers)


@app.get("/health")
async def health() -> JSONResponse:
    ready, details = await _backend_status()
    return JSONResponse(
        status_code=200 if ready else 503,
        content={"status": "healthy" if ready else "warming", **details},
    )


@app.get("/v1/models")
async def models_proxy() -> Response:
    async with httpx.AsyncClient(timeout=10.0) as client:
        response = await client.get(f"{BACKEND_URL}/v1/models")
    return Response(
        content=response.content,
        status_code=response.status_code,
        media_type=response.headers.get("content-type", "application/json"),
    )


@app.post("/v1/audio/speech")
async def speech_proxy(request: Request) -> Response:
    try:
        payload = await request.json()
    except ValueError as exc:
        raise HTTPException(status_code=400, detail="Request body must be valid JSON") from exc
    if not isinstance(payload, dict):
        raise HTTPException(status_code=400, detail="Request body must be a JSON object")

    try:
        response = await _generate(payload)
    except httpx.ConnectError as exc:
        raise HTTPException(status_code=503, detail="Model is still warming up") from exc
    except httpx.TimeoutException as exc:
        raise HTTPException(status_code=504, detail="Generation timed out") from exc
    forwarded_headers = {
        name: value
        for name, value in response.headers.items()
        if name.lower().startswith("x-") or name.lower() == "content-disposition"
    }
    return Response(
        content=response.content,
        status_code=response.status_code,
        media_type=response.headers.get("content-type", "application/octet-stream"),
        headers=forwarded_headers,
    )