# Fish Audio Drama 3 on SGLang Omni, as a Hugging Face Docker Space on one GPU
# (RTX Pro 6000 or H200).
#
# The base image is the one SGLang Omni's CI runs main on: CUDA, SGLang, and most
# Python dependencies, but not sglang-omni itself, which is installed at a pinned commit.
FROM hongccc/sglang-omni@sha256:ebe4239e29a764ee3a2806385c061c5fd438a26f01458e503d3822dcba5790df

ARG SGLANG_OMNI_COMMIT=22d4ace88f9547777b819a69d7691204bb9ce847

# The image moves slower than SGLang Omni (its SGLang predates the pinned commit's
# sglang==0.5.21). Install the commit's pins the image does not satisfy, as SGLang
# Omni's CI does before testing on this image (.github/scripts/prepare_omni_venv.sh).
COPY scripts/omni_missing_deps.py /tmp/omni_missing_deps.py
RUN python3 -c "import sys, urllib.request; urllib.request.urlretrieve(sys.argv[1], '/tmp/omni-pyproject.toml')" \
        "https://raw.githubusercontent.com/sgl-project/sglang-omni/${SGLANG_OMNI_COMMIT}/pyproject.toml" \
    && python3 /tmp/omni_missing_deps.py /tmp/omni-pyproject.toml > /tmp/omni-missing.txt \
    && echo "Pins the image does not satisfy:" && cat /tmp/omni-missing.txt \
    && if [ -s /tmp/omni-missing.txt ]; then xargs -d '\n' -n 1 python3 -m pip install --no-cache-dir < /tmp/omni-missing.txt; fi \
    && python3 /tmp/omni_missing_deps.py --overrides /tmp/omni-pyproject.toml \
        | xargs -r -d '\n' python3 -m pip install --no-cache-dir --no-deps \
    && python3 /tmp/omni_missing_deps.py --report /tmp/omni-pyproject.toml

RUN python3 -m pip install --no-cache-dir --no-deps \
        "sglang-omni @ https://github.com/sgl-project/sglang-omni/archive/${SGLANG_OMNI_COMMIT}.tar.gz" \
        python-multipart \
    && python3 -m pip install --no-cache-dir authlib itsdangerous "gradio[oauth]==6.28.0" \
    && python3 -c "import sglang_omni, whisper, av, librosa, httpx, fastapi, huggingface_hub; from huggingface_hub import attach_huggingface_oauth, parse_huggingface_oauth; print('sglang-omni', sglang_omni.__version__, 'huggingface_hub', huggingface_hub.__version__)"

# SGLang Omni declares each stage's factory in the model's config class and refuses
# to override it from the command line. Rebind the class name the Drama 3
# preprocessing factory builds to the Space's subclass, which adds described
# characters and directing history (drama3_space_ext.py, on PYTHONPATH at runtime).
RUN stages=$(python3 -c "import os, sglang_omni; print(os.path.join(os.path.dirname(sglang_omni.__file__), 'models', 'fishaudio_drama3', 'stages.py'))") \
    && grep -q "^def create_preprocessing_executor" "$stages" \
    && grep -q "preprocessor = Drama3SpeechPreprocessor(" "$stages" \
    && printf '\n\n# Drama 3 Space: described characters and directing history.\nfrom drama3_space_ext import Drama3SpacePreprocessor as Drama3SpeechPreprocessor  # noqa: E402\n' >> "$stages" \
    && tail -3 "$stages"

# Spaces run the container as uid 1000. Give it a writable home that holds the
# FlashInfer JIT cache the image built for Hopper, and stop the base entrypoint from
# cloning SGLang Omni's main branch at start.
RUN (id -u 1000 >/dev/null 2>&1 || useradd -m -u 1000 -s /bin/bash user) \
    && mkdir -p /home/user/app /home/user/.cache \
    && if [ -d /root/.cache/flashinfer ]; then cp -a /root/.cache/flashinfer /home/user/.cache/; fi \
    && chown -R 1000:1000 /home/user
ENV HOME=/home/user \
    USER=user \
    FLASHINFER_WORKSPACE_BASE=/home/user \
    SGLANG_OMNI_AUTO_CLONE=0 \
    PYTHONUNBUFFERED=1 \
    PYTHONPATH=/home/user/app \
    WHISPER_DIR=/home/user/.cache/whisper \
    DEMO_UNLIMITED_ORGS=fishaudio \
    USAGE_BUCKET_MOUNT_PATH=/data

USER 1000
WORKDIR /home/user/app

# Whisper large-v3-turbo transcribes reference clips (about 1.5 GB).
RUN python3 -c "import whisper; whisper._download(whisper._MODELS['turbo'], '/home/user/.cache/whisper', False)"

COPY --chown=1000:1000 . .

EXPOSE 7860
CMD ["python3", "server.py"]