# syntax=docker/dockerfile:1.6
#
# OmniVoice (k2-fsa/OmniVoice) — zero-shot, massively-multilingual (600+
# languages) voice-cloning + voice-design TTS, diffusion-LM architecture,
# Apache-2.0. Upstream ships a pip package + its own Gradio demo
# (`omnivoice-demo`); there's no official image, so we build a thin CUDA
# container around the pip package and run our OWN FastAPI wrapper (app.py):
# a batch OpenAI-style /v1/audio/speech plus a streaming /tts driven by the
# vendored buffer-ratchet scheduler (scheduler.py) for live chat consumers.

ARG CUDA_BASE=nvidia/cuda:12.8.0-cudnn-runtime-ubuntu22.04
FROM ${CUDA_BASE}

ENV DEBIAN_FRONTEND=noninteractive \
    PIP_ROOT_USER_ACTION=ignore \
    PYTHONUNBUFFERED=1 \
    PATH="/opt/venv/bin:${PATH}"

RUN apt-get update \
 && apt-get install -y --no-install-recommends \
        python3.10 python3.10-venv python3-pip \
        git ffmpeg libsndfile1 \
        ca-certificates wget \
 && rm -rf /var/lib/apt/lists/* \
 && python3.10 -m venv /opt/venv

# CUDA 12.8 torch wheels (per OmniVoice's documented install line).
RUN pip install --no-cache-dir \
        torch==2.8.0 torchaudio==2.8.0 \
        --index-url https://download.pytorch.org/whl/cu128

# OmniVoice from PyPI (+ huggingface_hub for the entrypoint weight pre-warm).
# Optional reproducible pin via the OMNIVOICE_VERSION build arg (empty=latest).
ARG OMNIVOICE_VERSION=
RUN pip install --no-cache-dir "omnivoice${OMNIVOICE_VERSION:+==${OMNIVOICE_VERSION}}" huggingface_hub

# Our asset-engine wrapper deps (FastAPI stack + soundfile for WAV encoding).
RUN pip install --no-cache-dir fastapi 'uvicorn[standard]' soundfile python-multipart

# Fail the build loudly if the runtime imports aren't satisfiable.
RUN python -c "import omnivoice, fastapi, soundfile, uvicorn; print('omnivoice wrapper deps OK')"

WORKDIR /app
COPY app.py /app/app.py
COPY scheduler.py /app/scheduler.py
COPY sanitize.py /app/sanitize.py
COPY entrypoint.sh /usr/local/bin/entrypoint.sh
RUN chmod +x /usr/local/bin/entrypoint.sh

# Fail the build if the wrapper (incl. the vendored scheduler + sanitizer) won't
# import. Model load is lazy (startup event), so this is a cheap CPU-only check.
RUN python -c "import app; print('omnivoice app import OK')"

EXPOSE 8001

# Our wrapper exposes /healthz (200 once model + >=1 voice are loaded). Generous
# start period: first boot pre-warms OmniVoice + Whisper ASR + clones every voice.
HEALTHCHECK --interval=30s --timeout=10s --start-period=1200s --retries=3 \
    CMD wget -q -O /dev/null http://127.0.0.1:8001/healthz || exit 1

ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8001"]