ARG BASE_IMAGE
FROM ${BASE_IMAGE}

ARG SOURCE_REF
ARG REFERENCE_WAV_PATH
ARG REFERENCE_WAV_SHA256
ARG PIP_VERSION
ARG PIP_INDEX_URL
ARG SETUPTOOLS_INSTALL_SPEC
ARG TORCH_VERSION_IDENTITY
ARG TORCH_LIBRARY_PATH
ARG TORCHAUDIO_VERSION_IDENTITY
ARG OPENAI_WHISPER_INSTALL_SPEC
ARG ONNXRUNTIME_GPU_INSTALL_SPEC
ARG ONNXRUNTIME_GPU_INDEX_URL

ENV DEBIAN_FRONTEND=noninteractive \
    PIP_INDEX_URL=${PIP_INDEX_URL} \
    PYTHONPATH=/opt/CosyVoice:/opt/CosyVoice/third_party/Matcha-TTS \
    PYTHONUNBUFFERED=1

RUN apt-get update \
    && apt-get install -y --no-install-recommends build-essential curl ffmpeg git git-lfs libsox-dev sox \
    && rm -rf /var/lib/apt/lists/* \
    && python -m pip install "pip==${PIP_VERSION}" "${SETUPTOOLS_INSTALL_SPEC}" wheel

COPY --from=cosyvoice_source . /opt/CosyVoice

RUN test "$(git -C /opt/CosyVoice rev-parse HEAD)" = "${SOURCE_REF}" \
    && test -z "$(git -C /opt/CosyVoice submodule status --recursive | sed -n '/^[+-U]/p')" \
    && printf '%s  %s\n' "${REFERENCE_WAV_SHA256}" "/opt/CosyVoice/${REFERENCE_WAV_PATH}" | sha256sum -c -

RUN sed -e '/^torch==/d' \
        -e '/^deepspeed==/d' \
        -e '/^tensorrt-cu12/d' \
        -e '/^openai-whisper==/d' \
        -e '/^onnxruntime-gpu==/d' \
        /opt/CosyVoice/requirements.txt > /tmp/runtime-requirements.txt \
    && python -m pip install "${OPENAI_WHISPER_INSTALL_SPEC}" --no-build-isolation

RUN python -m pip install --index-url "${ONNXRUNTIME_GPU_INDEX_URL}" "${ONNXRUNTIME_GPU_INSTALL_SPEC}"

RUN python -m pip install -r /tmp/runtime-requirements.txt \
    && test -d "${TORCH_LIBRARY_PATH}" \
    && python -c 'import sys, torch, torchaudio; expected_torch=sys.argv[1]; expected_audio=sys.argv[2]; assert torch.__version__.split("+")[0] == expected_torch, f"torch identity mismatch: expected={expected_torch} actual={torch.__version__}"; assert torchaudio.__version__.split("+")[0] == expected_audio, f"torchaudio identity mismatch: expected={expected_audio} actual={torchaudio.__version__}"; from cosyvoice.cli.cosyvoice import AutoModel' "${TORCH_VERSION_IDENTITY}" "${TORCHAUDIO_VERSION_IDENTITY}"

LABEL ai.unidesk.source.ref="${SOURCE_REF}"
LABEL ai.unidesk.torchaudio.version-identity="${TORCHAUDIO_VERSION_IDENTITY}"

WORKDIR /app
COPY api.py /app/api.py
COPY --chmod=755 model-init.sh /app/model-init.sh

EXPOSE 18327
CMD ["python", "/app/api.py"]
