# AuK gets its own image: it pins the PyTorch family to the 2.7 ABI line (torch>=2.7,<2.8 —
# every other backend here is on 2.8) and needs torchvision, which nothing else does.
# -runtime is enough: AuK is pure Python + prebuilt wheels, and the inference path is forced onto
# torch SDPA (`attn_backend = "torch"` in AukInfer), so flash-attn is never built or imported.
FROM nvidia/cuda:12.9.0-runtime-ubuntu24.04

# Avoid interactive prompts during package installation
ENV DEBIAN_FRONTEND=noninteractive

# System deps: libsndfile1 for soundfile I/O (writing generated wavs + decoding the dataset's
# prompt_audio column); ffmpeg for the librosa/soundfile decoders qwen_omni_utils uses to read
# the reference clip out of the ChatML turn; git for the pip VCS install of AuK itself.
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3 \
    python3-pip \
    python3-dev \
    git \
    ffmpeg \
    libsndfile1 \
    && rm -rf /var/lib/apt/lists/*

# Set Python alias (Ubuntu 24.04 ships Python 3.12; AuK requires >=3.10)
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

# PyTorch first (cu128 wheels for CUDA 12.8+/12.9 compat), pinned to the 2.7 line AuK's
# pyproject requires (torch>=2.7,<2.8 / torchaudio>=2.7,<2.8 / torchvision>=0.22,<0.23).
# Installing it up front stops `pip install auk` from resolving a CPU wheel for those pins.
# torchvision is a hard AuK dependency even though nothing in the TTS path uses it.
RUN pip install --no-cache-dir \
    torch==2.7.1 \
    torchaudio==2.7.1 \
    torchvision==0.22.1 \
    --index-url https://download.pytorch.org/whl/cu128

# transformers pinned to the release AuK validated against (4.57); Qwen2.5-Omni support landed in
# 4.52 and the Thinker's kwargs are a per-version detail, so this is not left to float.
RUN pip install --no-cache-dir "transformers==4.57.*" "accelerate>=0.33" "qwen-omni-utils>=0.0.9"

# AuK itself, pinned to a commit for reproducibility (not on PyPI). Core inference extra only —
# `[gradio]`/`[comfyui]` pull funasr/modelscope/tencentcloud, none of which this eval touches.
ARG AUK_COMMIT=bc84b758c7566dfce9bef0c49fba443456093b40
RUN pip install --no-cache-dir "git+https://github.com/Tencent-Hunyuan/AuK.git@${AUK_COMMIT}"

# Eval-loop deps. librosa is what qwen_omni_utils uses to read the reference wav named in the
# ChatML turn; soundfile writes the generated wavs and the prompt clips stage 3 scores SIM on.
#
# `audioread` is an UNDECLARED dependency of qwen-omni-utils 0.0.9: its v2_5/audio_process.py does
# a module-level `import audioread`, but its requires_dist is only av/librosa/packaging/pillow/
# requests. It used to arrive transitively through librosa, which dropped it in 0.11 (librosa is
# now on soundfile+soxr alone) — so on a fresh install the reference-audio path dies with
# ModuleNotFoundError inside CFMEdit.build_cond_inputs. Pin it here rather than holding librosa
# back at <0.11 for a transitive reason. Drop it if qwen-omni-utils ever declares it.
RUN pip install --no-cache-dir "datasets<4.0" tqdm soundfile librosa audioread

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