# Ubuntu 22.04 (Python 3.10) + CUDA 12.9, matching upstream's own environment: index-tts declares
# `requires-python = ">=3.10,<3.12"` and pins torch 2.8.* from the cu128 wheel index.
#
# The Python version is NOT cosmetic here. WeTextProcessing (index-tts's Linux text-normalization
# dependency) needs pynini, which ships prebuilt manylinux wheels only up to cp311 — on the
# Ubuntu 24.04 / Python 3.12 base the other backends use, pip would fall back to pynini's sdist
# and try to build OpenFst from source. 22.04's Python 3.10 gets the wheel.
#
# -runtime is enough: everything installs as prebuilt wheels, and BigVGAN's optional fused CUDA
# kernel is left off (see run_eval.py --use_cuda_kernel), so no nvcc is needed at run time either.
FROM nvidia/cuda:12.9.0-runtime-ubuntu22.04

# Avoid interactive prompts during package installation
ENV DEBIAN_FRONTEND=noninteractive

# Make apt resilient to slow/stalled mirrors: fail a stuck download fast and retry a few
# times instead of hanging for hours on a single package.
RUN printf 'Acquire::Retries "5";\nAcquire::http::Timeout "30";\nAcquire::https::Timeout "30";\n' \
    > /etc/apt/apt.conf.d/99-retries-timeout

# System deps: git to clone index-tts (it is not on PyPI); build-essential + python3-dev for the
# few deps without a cp310 wheel; ffmpeg/libsndfile1 for audio I/O (torchaudio/librosa/soundfile);
# libgl1 + libglib2.0-0 for opencv-python, which index-tts lists as a dependency and which
# otherwise fails at `import cv2` with "libGL.so.1: cannot open shared object file".
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3 \
    python3-pip \
    python3-dev \
    build-essential \
    git \
    ffmpeg \
    libsndfile1 \
    libgl1 \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

# Set Python alias (Ubuntu 22.04 ships Python 3.10, which is what upstream targets)
RUN ln -sf /usr/bin/python3 /usr/bin/python
ENV PIP_BREAK_SYSTEM_PACKAGES=1

WORKDIR /app

# Upgrade pip so it prefers prebuilt manylinux wheels over source builds.
RUN pip install --no-cache-dir --upgrade --ignore-installed pip "setuptools<81" wheel

# setuptools MUST stay <81 for the whole build, for the same reason cosyvoice3/Dockerfile pins it:
# index-tts depends on `openai-whisper`, an sdist whose setup.py does `import pkg_resources`, which
# setuptools 81 REMOVED. Pinning the outer setuptools is not enough — PEP 517 build isolation
# creates a fresh overlay env with the LATEST setuptools — so the constraint has to reach build
# environments too. PIP_CONSTRAINT does that today; PIP_BUILD_CONSTRAINT is the replacement pip
# 26.2+ requires, so set both and the build works either side of that change. Both persist into
# the runtime image on purpose: any pip install inside a job then resolves the same way.
RUN printf 'setuptools<81\n' > /etc/pip-constraints.txt
ENV PIP_CONSTRAINT=/etc/pip-constraints.txt \
    PIP_BUILD_CONSTRAINT=/etc/pip-constraints.txt

# Install the PyTorch ecosystem FIRST, from the cu128 index that index-tts's own `[tool.uv.index]`
# points at. Doing it up front means the `pip install -e` below sees torch==2.8.0 already
# satisfying its `torch==2.8.*` pin and keeps these CUDA wheels instead of resolving the default
# PyPI (CPU-ish) build over them.
RUN pip install --no-cache-dir \
    torch==2.8.0 \
    torchaudio==2.8.0 \
    --index-url https://download.pytorch.org/whl/cu128

# index-tts is NOT published on PyPI, so it is used as a source clone. Pinned to a commit rather
# than tracking `main`: this image defines the numbers the leaderboard reports, and an upstream
# change to the inference path would silently move them. Bump INDEXTTS_REF (and rebuild) to adopt
# a newer upstream.
ARG INDEXTTS_REF=ee40fa7d6c6b8a2c7f06105f9f1e65775b74868c
RUN git clone https://github.com/index-tts/index-tts.git /opt/index-tts \
    && git -C /opt/index-tts checkout --detach "${INDEXTTS_REF}"
ENV INDEXTTS_ROOT=/opt/index-tts

# Install index-tts and its dependencies. NO extras: `webui` (gradio), `deepspeed` and `accel`
# (flash-attn, which compiles for minutes) are all optional, and run_eval.py constructs IndexTTS2
# with use_deepspeed=False / use_accel=False / use_torch_compile=False, so none of them is on the
# inference path.
RUN pip install --no-cache-dir -e /opt/index-tts

# datasets + tqdm for the eval loop (soundfile/librosa/numpy come in via index-tts deps).
#
# datasets MUST stay <4, for the same reason cosyvoice3 pins it: from 4.0 on it decodes Audio()
# columns with torchcodec instead of soundfile, so an unpinned install dies on the eval set's
# `prompt_audio` with "To support decoding audio data, please install 'torchcodec'".
RUN pip install --no-cache-dir "datasets<4" tqdm soundfile

# NLTK corpora used by g2p-en (an index-tts dependency) for English G2P. Downloading them here
# keeps the job offline-ish and off nltk's servers at generation time. Tolerated failure: the
# corpus names have been renamed upstream before (`averaged_perceptron_tagger` ->
# `..._eng`), and this is a cache warm-up, not a requirement — nltk would fetch on demand.
RUN python3 -c "\
import nltk;\
[nltk.download(p, quiet=True) for p in \
 ('averaged_perceptron_tagger', 'averaged_perceptron_tagger_eng', 'cmudict', 'punkt', 'punkt_tab')]" \
    || echo 'NLTK corpus pre-download skipped (will be fetched on demand)'

# 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"]