# ZONOS3-Mini (Zyphra/ZONOS3-Mini) through the offline `zonos3_sgl.tts.TTSLLM` API of
# https://github.com/Zyphra/ZONOS3-Mini (Mini-SGLang server package), set up as its README says:
# the CUDA 12.8 toolkit (the LM JIT-compiles its kernels at first start, so the devel image with nvcc),
# Python 3.12 and `uv sync --locked --all-packages`. The weights download from the public model repo
# on first use (no token needed).
FROM nvidia/cuda:12.8.1-devel-ubuntu24.04

ENV DEBIAN_FRONTEND=noninteractive

# git: the repo + its FuseAlign submodule (a workspace path dependency). build-essential: host
# compiler for the JIT kernels (Triton / TileLang / SGL). ffmpeg: audio decoding of the eval
# datasets (torchcodec) and of the mp3 default voices.
RUN apt-get update && apt-get install -y --no-install-recommends \
    git ca-certificates build-essential ffmpeg libsndfile1 \
    && rm -rf /var/lib/apt/lists/*

# uv, with its managed Python 3.12 outside /root so any user can run it.
COPY --from=ghcr.io/astral-sh/uv:0.9.7 /uv /uvx /usr/local/bin/
ENV UV_PYTHON_INSTALL_DIR=/opt/uv-python UV_LINK_MODE=copy UV_NO_CACHE=1

# The repo at a pinned commit. Only the FuseAlign submodule: ggml is for the C++ server, unused here.
ARG ZONOS3_COMMIT=87d1bacf5be9d45b1766b14bff0ce5cea5941428
RUN git clone https://github.com/Zyphra/ZONOS3-Mini.git /opt/ZONOS3-Mini \
    && cd /opt/ZONOS3-Mini && git checkout ${ZONOS3_COMMIT} \
    && git submodule update --init --depth 1 third_party/fusealign-multilingual

WORKDIR /opt/ZONOS3-Mini
# The locked environment (torch 2.9.1 cu128, sgl-kernel, flashinfer, tilelang, ...).
RUN uv sync --locked --all-packages --no-dev

# Eval-loop deps the workspace lacks, constrained to the lockfile so nothing locked moves:
# datasets (5.0.1 is the newest that accepts the locked huggingface_hub 0.36) and torchcodec, its
# audio decoder (0.8.x is the release built against torch 2.9).
RUN uv export --locked --no-hashes --all-packages --no-emit-workspace --no-dev -q \
        | grep -v '^\s*#\|^-e \|^\.' > /tmp/lock-constraints.txt \
    && uv pip install --python .venv/bin/python -c /tmp/lock-constraints.txt \
        "datasets==5.0.1" "torchcodec==0.8.1" \
    && rm /tmp/lock-constraints.txt

# The workspace venv is the default Python: submit_jobs.sh runs `source servers/sgl/scripts/env.sh`
# (CUDA toolchain, kernel selection per GPU) then `python run_eval.py`.
ENV VIRTUAL_ENV=/opt/ZONOS3-Mini/.venv PATH=/opt/ZONOS3-Mini/.venv/bin:$PATH

# Fail the BUILD, not a job, on a broken install. The builder has no GPU driver, and sgl_kernel
# (imported by zonos3_sgl.tts) needs libcuda.so.1 at import, so check the GPU-free parts: the
# eval deps, the model package, that zonos3_sgl.tts is installed, and the default voice.
RUN python -c "\
import importlib.util, torch, torchcodec, datasets, soundfile as sf, zonos3, zonos3.languages; \
assert importlib.util.find_spec('zonos3_sgl.tts') and importlib.util.find_spec('sgl_kernel'); \
w, sr = sf.read('assets/default_voices/English/Savannah_en-American_F_1.mp3'); \
print('torch', torch.__version__, torch.version.cuda, '| torchcodec', torchcodec.__version__, \
      '| datasets', datasets.__version__, '| voice', w.shape, sr)"

WORKDIR /app
# The eval scripts are injected at runtime by submit_jobs.sh.
COPY . /app

ENTRYPOINT ["bash"]

# Keep-alive CMD so the Space runtime stays healthy; `docker run` overrides it.
EXPOSE 7860
CMD ["-c", "python3 -m http.server 7860"]