# -runtime is enough: neutts ships pure-Python + prebuilt wheels (no CUDA extensions to compile).
FROM nvidia/cuda:12.9.0-runtime-ubuntu24.04

# Avoid interactive prompts during package installation
ENV DEBIAN_FRONTEND=noninteractive

# System deps: espeak-ng is the phonemizer backend used by neutts (the pip package bundles
# a copy but falls back to the system library); ffmpeg + libsndfile1 for audio I/O.
# build-essential + cmake are BUILD-time only: they compile the CPU llama.cpp backend below.
# (They also satisfy Triton's runtime JIT if a future change puts this backend on a
# torch.compile path — the same gap that broke voxcpm2's streaming probe.)
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3 \
    python3-pip \
    python3-dev \
    build-essential \
    cmake \
    git \
    ffmpeg \
    libsndfile1 \
    espeak-ng \
    && rm -rf /var/lib/apt/lists/*

# Set Python alias (Ubuntu 24.04 ships Python 3.12)
RUN ln -sf /usr/bin/python3 /usr/bin/python

# Allow pip to install packages system-wide in the container (PEP 668)
ENV PIP_BREAK_SYSTEM_PACKAGES=1

WORKDIR /app

# Upgrade pip so it prefers prebuilt manylinux wheels. --ignore-installed is required on
# Ubuntu 24.04: the Debian-installed pip has no RECORD file and cannot be uninstalled.
RUN pip install --no-cache-dir --upgrade --ignore-installed pip setuptools wheel

# Install PyTorch ecosystem (cu128 wheels for CUDA 12.8+/12.9 compat). neutts needs torch>=2.8.0.
RUN pip install --no-cache-dir \
    torch==2.8.0 \
    torchaudio==2.8.0 \
    --index-url https://download.pytorch.org/whl/cu128

# neutts (NeuTTS Nano/Air inference package; brings transformers~=4.56, neucodec, phonemizer,
# librosa, soundfile, resemble-perth). Torch GPU path only — the llama-cpp (GGUF) and ONNX
# extras are not needed. This container does TTS generation ONLY — ASR/WER scoring runs in
# the separate `tts-eval` image.
RUN pip install --no-cache-dir neutts

# llama-cpp-python provides the GGUF backbone. It is what makes NeuTTS.infer_stream() work at all:
# on the torch backbone it raises NotImplementedError, so the quantized path is the only way to
# measure a real time-to-first-audio for this model (see run_eval.py's TTFA probe).
# The CUDA variant comes from a prebuilt wheel (compiling it would need nvcc, absent from this
# -runtime base); the CPU variant is compiled below, which needs only a host toolchain.
#
# BOTH variants are installed because the two link differently and this one image runs on both
# GPU and CPU flavors: the CUDA build's libllama.so links libcuda.so.1, which does not exist on
# a CPU-only machine, so importing it there dies with
#   RuntimeError: Failed to load shared library ... libcuda.so.1: cannot open shared object file
# The CUDA build is the default (site-packages); the CPU build sits in /opt/llama-cpp-cpu and
# run_eval.py puts it first on sys.path when --device=cpu. --no-deps on the second install
# because its pure-python dependencies are already satisfied system-wide.
RUN pip install --no-cache-dir llama-cpp-python \
        --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu124
# The CPU build is COMPILED here rather than taken from the prebuilt cpu wheel index. That index
# does carry this same version for Linux, but its wheels are built for portability with SIMD off;
# measured on cpu-upgrade the prebuilt one gave RTFx 0.07 and a 61 s TTFA — ~88x slower than this
# image's CUDA build and ~40x off tts-bench's CPU figures for the same model. That number
# describes the wheel, not NeuTTS.
#
# Pinned to the SAME version as the CUDA wheel above so a CPU-vs-GPU comparison differs only in
# the backend.
#
# GGML_NATIVE is deliberately OFF: it tunes for the BUILD machine, and the Space builder's CPU
# need not match the job's — an AVX-512 build would die with SIGILL on a host without it. AVX2 +
# FMA are safe on any x86-64 server CPU since ~2013 and recover most of the SIMD speedup.
#
# CMAKE_BUILD_PARALLEL_LEVEL is deliberately SMALL: this trades build time for peak memory, and
# memory is the binding constraint on the HF Space builder.
#
# Measured here (cache busted each time): 2 jobs -> 206 s; the default (all cores) -> 31 s. So
# parallelism genuinely does help — but the builder reports the HOST's core count while capping
# memory, so letting it use every core spawns dozens of g++ and the build is OOM-killed. That
# failure is silent: it ends on "Building wheel ... still running..." with no error, exactly as
# flash-attn did in breeze-tts/Dockerfile at MAX_JOBS=8. 3.5 minutes is a cheap insurance premium.
# Raise LLAMA_BUILD_JOBS if your builder has memory to spare.
#
# The LLAMA_BUILD_* switches are kept for tidiness; measured, they change neither time nor size.
ARG LLAMA_CPP_VERSION=0.3.19
ARG LLAMA_BUILD_JOBS=2
RUN CMAKE_BUILD_PARALLEL_LEVEL="${LLAMA_BUILD_JOBS}" \
    CMAKE_ARGS="-DGGML_NATIVE=OFF -DGGML_AVX=ON -DGGML_AVX2=ON -DGGML_FMA=ON -DGGML_F16C=ON \
                -DGGML_CUDA=OFF -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF \
                -DLLAMA_BUILD_SERVER=OFF -DLLAMA_BUILD_TOOLS=OFF -DLLAMA_BUILD_COMMON=OFF" \
    FORCE_CMAKE=1 \
    pip install --no-cache-dir --no-deps --target=/opt/llama-cpp-cpu \
        "llama-cpp-python==${LLAMA_CPP_VERSION}"

# gguf is needed to read the array metadata llama-cpp drops (_read_gguf_array_meta).
RUN pip install --no-cache-dir gguf
# neucodec declares torchao>=0.12 and pip resolves to the latest torchao, which is built
# against torch >= 2.11 and warns "Skipping import of cpp extensions" under torch 2.8.
# Pin the torchao release paired with torch 2.8 (0.13.0, Sep 2025) instead.
RUN pip install --no-cache-dir "torchao==0.13.0"

# datasets + tqdm for the eval loop (soundfile comes in via neutts deps).
# torchcodec: required by `datasets` to decode the Audio column (the cloning reference
# clips); 0.7.x is the release paired with torch 2.8 and uses the system ffmpeg.
RUN pip install --no-cache-dir datasets tqdm soundfile "torchcodec==0.7.*"

# Weights are NOT baked in: the `neutts` package selects the codec class by exact repo id
# ("neuphonic/neucodec"), so both backbone (~0.5 GB) and codec are pulled from the Hub into
# the HF cache at first run (both ungated; run_eval.sh bind-mounts the host HF cache).

# Copy the full repository
COPY . /app

# Default entrypoint
ENTRYPOINT ["bash"]

# Keep-alive CMD so the Space runtime stays healthy. HF Jobs and `docker run`
# override this with their own command (e.g. run_eval.sh).
EXPOSE 7860
CMD ["-c", "python3 -m http.server 7860"]