voice_copy/app.py

"""
voice_copy — simple web UI for zero-shot voice cloning.

Run:
    python app.py            # then open http://localhost:8000
    # or: uvicorn app:app --host 0.0.0.0 --port 8000

Env:
    VOICE_COPY_MODEL_SIZE   "1.7B" (default) or "0.6B"
    VOICE_COPY_DEVICE       force "cuda" / "cpu" / "mps" (default: auto)
    VOICE_COPY_HOST         bind host (default 0.0.0.0)
    VOICE_COPY_PORT         bind port (default 8000)
"""

from __future__ import annotations

import io
import logging
import os
import tempfile
from pathlib import Path

import numpy as np
import soundfile as sf
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.concurrency import run_in_threadpool
from fastapi.responses import HTMLResponse, Response

from engine import HF_MODEL_MAP, LANGUAGE_CODE_TO_NAME, VoiceCloner

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger("voice_copy.app")

BASE = Path(__file__).resolve().parent
STATIC = BASE / "static"
TMP = BASE / "tmp"
TMP.mkdir(exist_ok=True)

MODEL_SIZE = os.environ.get("VOICE_COPY_MODEL_SIZE", "1.7B")
DEVICE = os.environ.get("VOICE_COPY_DEVICE") or None

cloner = VoiceCloner(model_size=MODEL_SIZE, device=DEVICE)
app = FastAPI(title="voice_copy", description="Zero-shot voice cloning UI")


def _decode_to_wav(raw_bytes: bytes, filename: str) -> str:
    """Decode uploaded audio (wav/flac/mp3/m4a/ogg/webm) into a mono wav on disk.

    Tries soundfile first (wav/flac/ogg — no ffmpeg needed); falls back to
    librosa/audioread (needs ffmpeg for mp3/m4a/webm from the mic recorder).
    """
    suffix = Path(filename or "ref").suffix or ".bin"
    src = tempfile.NamedTemporaryFile(dir=TMP, suffix=suffix, delete=False)
    src.write(raw_bytes)
    src.close()

    try:
        try:
            audio, sr = sf.read(src.name, dtype="float32", always_2d=False)
            if audio.ndim > 1:  # stereo -> mono
                audio = audio.mean(axis=1)
        except Exception:
            import librosa

            audio, sr = librosa.load(src.name, sr=None, mono=True)
    finally:
        try:
            os.unlink(src.name)
        except OSError:
            pass

    wav = tempfile.NamedTemporaryFile(dir=TMP, suffix=".wav", delete=False)
    wav.close()
    sf.write(wav.name, np.asarray(audio, dtype=np.float32), int(sr), format="WAV")
    return wav.name


@app.get("/", response_class=HTMLResponse)
def index() -> str:
    return (STATIC / "index.html").read_text(encoding="utf-8")


@app.get("/api/health")
def health() -> dict:
    return {
        "device": cloner.device,
        "model_size": cloner.model_size,
        "model_repo": HF_MODEL_MAP[cloner.model_size],
        "model_loaded": cloner.loaded,
        "languages": list(LANGUAGE_CODE_TO_NAME.keys()),
    }


@app.post("/api/clone")
async def clone(
    ref_audio: UploadFile = File(...),
    ref_text: str = Form(...),
    text: str = Form(...),
    language: str = Form("auto"),
    instruct: str = Form(""),
    seed: str = Form(""),
) -> Response:
    if not text.strip():
        raise HTTPException(400, "합성할 텍스트(text)를 입력하세요.")
    if not ref_text.strip():
        raise HTTPException(400, "참조 오디오의 전사(ref_text)를 입력하세요.")

    raw = await ref_audio.read()
    if not raw:
        raise HTTPException(400, "참조 오디오 파일이 비어 있습니다.")

    seed_int = None
    if seed.strip():
        try:
            seed_int = int(seed)
        except ValueError:
            raise HTTPException(400, "seed는 정수여야 합니다.")

    try:
        wav_path = await run_in_threadpool(_decode_to_wav, raw, ref_audio.filename or "ref.wav")
    except Exception as exc:  # noqa: BLE001
        logger.exception("audio decode failed")
        raise HTTPException(
            400,
            f"오디오 디코딩 실패: {exc}. wav/flac는 그대로 되고, mp3/m4a/webm은 ffmpeg 설치가 필요합니다.",
        )

    try:
        audio, sr = await run_in_threadpool(
            cloner.clone, wav_path, ref_text, text, language, instruct or None, seed_int
        )
    except Exception as exc:  # noqa: BLE001
        logger.exception("clone failed")
        raise HTTPException(500, f"생성 실패: {exc}")
    finally:
        try:
            os.unlink(wav_path)
        except OSError:
            pass

    buf = io.BytesIO()
    sf.write(buf, audio, sr, format="WAV")
    buf.seek(0)
    return Response(
        content=buf.read(),
        media_type="audio/wav",
        headers={"Content-Disposition": 'inline; filename="cloned.wav"'},
    )


if __name__ == "__main__":
    import uvicorn

    host = os.environ.get("VOICE_COPY_HOST", "0.0.0.0")
    port = int(os.environ.get("VOICE_COPY_PORT", "8000"))
    uvicorn.run(app, host=host, port=port)