# Custom Dockerfile for ACE-Step 1.5.
#
# Mirrors upstream's Dockerfile structure, but fixes a CUDA-version
# mismatch that crashloops the upstream image as of April 2026:
#   * upstream's requirements.txt lists `torch torchvision torchaudio`
#     with no version pins;
#   * upstream's pip install uses `--extra-index-url cu126`, which is
#     a FALLBACK only — pypi default wins for resolution;
#   * pypi-default torch is now cu13, so torch installs cu13 + the
#     cu126 fallback only kicks in for torchvision/torchaudio →
#     `RuntimeError: Detected that PyTorch and torchvision were compiled
#     with different CUDA major versions`.
#
# Fix: install torch/torchvision/torchaudio FIRST from the cu126 index
# (forced via --index-url, not --extra-index-url). Then `pip install
# -r requirements.txt` sees they're already satisfied and leaves them
# alone.
#
# Also: command is `python3 infer-api.py` (REST), not `gui.py` (Gradio)
# — see compose.yaml command override; CMD here is the same default
# so the image works standalone too.
FROM nvidia/cuda:12.6.0-runtime-ubuntu22.04 AS base

ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    HF_HUB_ENABLE_HF_TRANSFER=1 \
    DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y --no-install-recommends \
        python3.10 \
        python3-pip \
        python3-venv \
        python3-dev \
        build-essential \
        git \
        curl \
        ca-certificates \
        ffmpeg \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sf /usr/bin/python3 /usr/bin/python
# ffmpeg is required by torchcodec at runtime — torchcodec dlopens
# libavcodec/libavformat. Without it the WAV save step fails with
# "Could not load libtorchcodec".

RUN python -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

WORKDIR /app

# Clone upstream. Bake the SHA into a layer-cache key so a different
# SHA invalidates everything below.
ARG ACE_STEP_REF=main
RUN git clone https://github.com/ace-step/ACE-Step.git . \
    && git checkout ${ACE_STEP_REF} \
    && echo "ace-step ref: $(git rev-parse HEAD)"

# Pre-install torch/torchvision/torchaudio from the cu126 index — this
# satisfies the unpinned entries in requirements.txt so the next pip
# install doesn't re-resolve them from pypi default (cu13).
RUN pip install --no-cache-dir --upgrade pip \
    && pip install --no-cache-dir \
        torch torchvision torchaudio \
        --index-url https://download.pytorch.org/whl/cu126 \
    && pip install --no-cache-dir hf_transfer peft \
    && pip install --no-cache-dir -r requirements.txt \
    && pip install --no-cache-dir . \
    && pip install --no-cache-dir torchcodec
# torchcodec — required by torchaudio's save_with_torchcodec (the save
# path ACE-Step uses on output). Not pulled in by upstream's
# requirements.txt; without it, /generate runs to completion and then
# 500s at the WAV write step.

# Replace upstream's infer-api.py with our patched copy. Upstream's
# version builds an 18-arg positional tuple but the pipeline expects
# 24 — see infer-api.py header comment for the fix.
COPY infer-api.py /app/infer-api.py

EXPOSE 8000
CMD ["python3", "infer-api.py"]