# BlueTTS 2.5 ONNX + NVIDIA GPU for RunPod Serverless (no gradio).
# Build from repo root:
#   docker buildx build --platform linux/amd64 -f docker/Dockerfile \
#     -t maxme123/blue-v2:serverless-v8 --push .
#
# Russian is included: RUAccent + RUPhon and their ~1 GB of weights are baked in
# below so a cold worker does not download them on its first Russian job. They
# need torch, which is installed from the CPU wheel index — the front end runs on
# CPU by design (BLUE_RU_G2P_DEVICE), and the CUDA torch build would add ~2 GB of
# libraries that duplicate what the base image and onnxruntime-gpu already carry.
# Build with --build-arg WITH_RUSSIAN=0 for a torch-free image without Russian.

ARG CUDA_IMAGE=nvidia/cuda:12.6.3-cudnn-runtime-ubuntu24.04
FROM ${CUDA_IMAGE}

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    VIRTUAL_ENV=/opt/venv \
    PATH="/opt/venv/bin:$PATH" \
    PYTHONPATH=/app \
    BLUE_DOCKER=1 \
    BLUE_ONNX_GPU=1 \
    BLUE_SERVE_MODE=serverless \
    BLUE_SKIP_GRADIO=1 \
    GRADIO_ANALYTICS_ENABLED=False \
    HF_HUB_DISABLE_TELEMETRY=1 \
    ORT_NUM_THREADS=4 \
    MAX_CONCURRENCY=4

RUN apt-get update && apt-get install -y --no-install-recommends \
        python3 \
        python3-pip \
        python3-venv \
        python3-dev \
        git \
        ffmpeg \
        espeak-ng \
        libespeak-ng1 \
        libsndfile1 \
        ca-certificates \
        curl \
    && python3 --version \
    && ffmpeg -version | head -1 \
    && python3 -m venv /opt/venv \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

ARG WITH_RUSSIAN=1

COPY docker/requirements-serverless.txt /app/docker/requirements-serverless.txt
COPY docker/requirements-russian.txt /app/docker/requirements-russian.txt
# Order matters, and not for readability: ruphon depends on the CPU `onnxruntime`,
# which unpacks over the same directory as `onnxruntime-gpu`. Whichever pip writes
# last wins, so the Russian deps go in first and the serverless set — which pins
# onnxruntime-gpu — goes in last. The assert below is what turns a future flip
# into a failed build instead of a silently CPU-bound image.
RUN pip install --upgrade pip setuptools wheel \
    && if [ "$WITH_RUSSIAN" = "1" ]; then \
           pip install --index-url https://download.pytorch.org/whl/cpu torch \
           && pip install -r /app/docker/requirements-russian.txt; \
       fi \
    && pip install -r /app/docker/requirements-serverless.txt \
    && python -c "import onnxruntime as ort, sys; \
p = ort.get_available_providers(); \
print('ORT', ort.__version__, p); \
sys.exit('CUDAExecutionProvider missing — the CPU onnxruntime wheel clobbered the GPU one; check install order above') \
if 'CUDAExecutionProvider' not in p else None" \
    && if [ "$WITH_RUSSIAN" != "1" ]; then \
           python -c "import torch" 2>/dev/null && echo 'ERROR: torch should not be installed' && exit 1 || echo 'OK: no torch'; \
       fi

COPY app.py download_models.py tts.json vocab.json /app/
COPY g2p /app/g2p
COPY voices /app/voices
COPY onnx_slim /app/onnx_slim
COPY renikud.onnx /app/renikud.onnx
COPY docker/entrypoint.sh /app/docker/entrypoint.sh
COPY docker/rp_handler.py /app/docker/rp_handler.py

# Fail the build, not the first job, if the bundle is incomplete. `stats.npz` is
# what the vocoder input is denormalized with and `uncond.npz` is what makes
# cfg_scale do anything, so a v2 bundle baked in here would be a silent
# quality regression rather than a crash.
#
# The freshness markers are (re)written here rather than copied, so the baked-in
# weights are trusted at startup. Without them every cold start re-downloads
# renikud.onnx — 310 MB on the critical path of each new serverless worker.
RUN chmod +x /app/docker/entrypoint.sh \
    && for f in vocoder.onnx text_encoder.onnx vector_estimator.onnx \
                duration_predictor_style.onnx stats.npz uncond.npz vocab.json; do \
           test -s "/app/onnx_slim/$f" || { echo "MISSING onnx_slim/$f"; exit 1; }; \
       done \
    && test -s /app/renikud.onnx \
    && python -c "import sys, os; \
from download_models import VOICE_CATALOG, BLUE_REPO, ONNX_BUNDLE_STAMP, RENIKUD_REPO, RENIKUD_STAMP; \
missing=[m['file'] for m in VOICE_CATALOG.values() if not os.path.exists('/app/voices/'+m['file'])]; \
sys.exit('MISSING voices: %s' % missing) if missing else print('voices OK:', list(VOICE_CATALOG)); \
open('/app/.renikud_repo','w').write(RENIKUD_REPO+chr(10)+RENIKUD_STAMP+chr(10)); \
open('/app/onnx_slim/.repo_id','w').write(BLUE_REPO+chr(10)+ONNX_BUNDLE_STAMP+chr(10)); \
open('/app/voices/.repo_id','w').write(BLUE_REPO+chr(10)+ONNX_BUNDLE_STAMP+chr(10)); \
print('markers written')"

# Bake the Russian weights in. RUAccent's dictionary + omograph model and RUPhon's
# graph are ~1 GB fetched on first use, which would otherwise land on the first
# Russian job of every cold worker. Phonemizing one sentence is what triggers the
# download, and it doubles as a build-time check that the front end actually loads
# (the transformers/tokenizer patch in g2p/russian_g2p.py is version-sensitive).
RUN if [ "$WITH_RUSSIAN" = "1" ]; then \
        python -c "from g2p.russian_g2p import phonemize_russian; \
print('ru warm:', phonemize_russian('Взгляну́в на Санкт-Петербург, я всё понял.'))"; \
    fi

# No HEALTHCHECK — RunPod serverless uses the runpod SDK job loop, not HTTP probes.
CMD ["/app/docker/entrypoint.sh"]