# Breeze TTS 2 gets its own image: the inference repo hard-pins torch==2.9.1 /
# transformers==4.57.3 / qwen-tts==0.1.1, and the checkpoint requires flash-attn (see below).
#
# Base + Python version are load-bearing here. Dao-AILab publishes prebuilt flash-attn wheels for
# torch 2.9 in exactly ONE flavour — cu12 / cxx11abiTRUE / **cp312** / x86_64 — so the image must
# be Python 3.12 on x86_64. Ubuntu 24.04 ships 3.12, and a `-runtime` base is enough because
# nothing is compiled from source (an earlier revision built flash-attn here with MAX_JOBS=8 and
# was OOMKilled by the HF Space builder — nvcc peaks at several GB per parallel job).
FROM nvidia/cuda:12.9.0-runtime-ubuntu24.04

# Pin the inference repo to a commit: `models/` registers the `breeze` architecture that stock
# transformers does not know, so its API is what run_eval.py imports — a moving `main` would
# silently change the eval. Bump deliberately.
ARG BREEZE_REPO_URL=https://github.com/breezeblue-ai/breeze-tts.git
ARG BREEZE_REPO_SHA=ca632ce6c4d05f7985da4eab29b1a5d445b43f7b

# Prebuilt flash-attn wheel matching torch 2.9 + cu12 + cp312 + cxx11abi=TRUE. TRUE is not a
# choice: every torch 2.9 manylinux_2_28 wheel is built with _GLIBCXX_USE_CXX11_ABI=1, which is
# why upstream ships no FALSE variant for torch 2.9 at all.
ARG FLASH_ATTN_WHEEL=https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whl

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    TOKENIZERS_PARALLELISM=false

# System deps: git to clone the inference repo; libsndfile1/ffmpeg for audio I/O (writing
# generated wavs and decoding the dataset's prompt_audio column); sox is a runtime dependency of
# the `sox` Python package that qwen-tts pulls in. No build-essential/ninja/nvcc — see the base
# image note above: every Python dependency here installs from a wheel.
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3 \
    python3-pip \
    python3-dev \
    ca-certificates \
    curl \
    ffmpeg \
    git \
    libsndfile1 \
    sox \
    && rm -rf /var/lib/apt/lists/*

# Set Python alias (Ubuntu 24.04 ships Python 3.12 — required by the flash-attn wheel above)
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 first, from the cu128 index, so the vendor requirements below see torch==2.9.1
# as already satisfied and do not pull a different build from PyPI. (`==2.9.1` matches the local
# version `2.9.1+cu128` under PEP 440, so the pin is honoured.)
RUN pip install --no-cache-dir \
    torch==2.9.1 \
    torchaudio==2.9.1 \
    --index-url https://download.pytorch.org/whl/cu128

# Clone the inference repo at the pinned commit. It is NOT pip-installed — `models` and
# `breeze_infer` are plain packages imported by path (run_eval.py prepends BREEZE_REPO to
# sys.path), and `configs/fast.json` is read from the same tree by the fast-path warmup.
RUN git clone "${BREEZE_REPO_URL}" /opt/breeze-tts \
    && cd /opt/breeze-tts && git checkout --quiet "${BREEZE_REPO_SHA}"
ENV BREEZE_REPO=/opt/breeze-tts

# The vendor's exact pinned dependency set (torch/torchaudio already satisfied above).
RUN pip install --no-cache-dir -r /opt/breeze-tts/requirements.txt

# flash-attn is NOT optional here, even though run_eval.py runs the backbone with
# attn_implementation=eager: the checkpoint's text_encoder_config pins
# `preferred_attn_implementation: flash_attention_2`, so the T5Gemma2 text encoder imports
# flash_attn on every load. Installed from the prebuilt wheel rather than compiled — the wheel
# carries kernels for every supported arch, so unlike a source build this image is not tied to
# one GPU generation. --no-deps because its metadata would otherwise re-resolve torch.
RUN pip install --no-cache-dir --no-deps "${FLASH_ATTN_WHEEL}"

# datasets + tqdm + soundfile + librosa for the eval loop. datasets is pinned <4.0 deliberately:
# 4.0 switched the Audio decoder to torchcodec, and run_eval.py reads `prompt_audio` in the 3.x
# dict form (`["array"]` / `["sampling_rate"]`). librosa is REQUIRED, not optional — datasets 3.x's
# `Audio.decode_example` imports BOTH librosa and soundfile and raises if either is missing
# (qwen-tts happens to pull librosa in too, but do not rely on a transitive dep for that).
RUN pip install --no-cache-dir "datasets<4.0" tqdm soundfile librosa huggingface_hub

# Fail the build rather than an 8-hour job if the pinned stack did not actually land. This checks
# the exact torch/transformers/qwen-tts/flash-attn versions and imports flash_attn.
RUN python /opt/breeze-tts/docker/smoke_check.py

# NOTE: the weights are deliberately NOT baked into this image, unlike most backends here.
# BreezeBlue/Breeze-TTS-2 is ungated, but its weights carry the BreezeBlue Research and
# Non-Commercial License, whose §4(f) requires that "each recipient independently accepts this
# Agreement before accessing or using the Model Materials". This image is published as a PUBLIC
# Space for HF Jobs to pull, so baking the weights in would hand them to pullers who never
# accepted it. run_eval.py snapshot_downloads them at runtime instead (into the HF cache, so a
# mounted cache makes local reruns a no-op); HF_TOKEN is passed through but is not required.

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