FROM python:3.12-slim-bookworm

ARG YUE2_COMMIT=9c6c4b349be978b06a9d0d958471a07a6cdeff4d
ARG YUE2_WEBUI_COMMIT=8fc05609bde5dcd345d7c6d57fff3da0839164c5

RUN apt-get update \
    && apt-get install -y --no-install-recommends ca-certificates ffmpeg git libsndfile1 \
    && rm -rf /var/lib/apt/lists/*

RUN git clone https://github.com/multimodal-art-projection/YuE.git /opt/yue2 \
    && git -C /opt/yue2 checkout "${YUE2_COMMIT}" \
    && python -m pip install --no-cache-dir /opt/yue2

# The community console is intentionally layered onto the already verified,
# pinned YuE2 runtime.  We use only its webui/ directory; its bundled YuE2 fork
# and Windows installers never become part of the Athena image.
RUN git clone https://github.com/Ladypoly/YuE2_WebUI.git /tmp/yue2-webui \
    && git -C /tmp/yue2-webui checkout "${YUE2_WEBUI_COMMIT}" \
    && cp -a /tmp/yue2-webui/webui /opt/yue2/webui \
    && cp /tmp/yue2-webui/LICENSE /opt/yue2/YuE2_WebUI-LICENSE \
    && python -m pip install --no-cache-dir -r /opt/yue2/webui/requirements.txt \
    && rm -rf /tmp/yue2-webui

# SheetSage2 converts an uploaded recording into the ABC score YuE2 consumes.
# Keep its pinned Transformers/Numpy stack in a small virtual environment, but
# reuse the image's Blackwell-capable torch 2.10 + CUDA 12.8 installation.
# The upstream torch 2.8+cu126 recipe cannot execute RTX 50-series kernels.
RUN python -m venv --system-site-packages /opt/yue2/.venv-sheetsage2 \
    && /opt/yue2/.venv-sheetsage2/bin/pip install --no-cache-dir \
       "torchaudio==2.10.0" --index-url https://download.pytorch.org/whl/cu128 \
    && /opt/yue2/.venv-sheetsage2/bin/pip install --no-cache-dir \
       "transformers==4.45.2" "huggingface-hub==0.36.0" \
       "safetensors==0.5.3" "numpy==1.26.4" "scipy==1.13.1" \
       "mir_eval==0.8.2" "pretty_midi==0.2.10" "mido==1.3.3" \
       "setuptools==78.1.1"

COPY community-webui-patches /tmp/community-webui-patches
RUN python /tmp/community-webui-patches/patch_webui.py /opt/yue2/webui \
    && rm -rf /tmp/community-webui-patches

COPY ui /opt/yue2-playground

WORKDIR /workspace
ENV PYTHONUNBUFFERED=1 \
    YUE2_KIT=/workspace

ENTRYPOINT ["yue2"]