# Stable Audio Open 1.0 inference image.
# pytorch/pytorch base ships torch + cuda + cudnn already linked, so
# we only layer the diffusers stack + a libsndfile for soundfile + the
# fastapi shim. Smaller and faster to build than starting from
# nvidia/cuda and pip-installing torch ourselves.
FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime AS base

ENV PYTHONUNBUFFERED=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    PIP_NO_CACHE_DIR=1 \
    PIP_DISABLE_PIP_VERSION_CHECK=1 \
    HF_HOME=/app/hf_cache

# libsndfile1 is the C lib soundfile binds to. Without it the pip
# install of soundfile succeeds but `import soundfile` fails at
# runtime with OSError: cannot find libsndfile.
RUN apt-get update && apt-get install -y --no-install-recommends \
        libsndfile1 \
    && rm -rf /var/lib/apt/lists/*

# protobuf + sentencepiece are pulled in by the T5 text encoder
# (Stable Audio Open uses google/t5-base-cb under the hood).
# accelerate gates the .to(device) fast path for diffusers.
# torchsde is required by CosineDPMSolverMultistepScheduler — diffusers
# doesn't pull it as a hard dep; without it, pipeline init fails with
# "CosineDPMSolverMultistepScheduler requires the torchsde library".
RUN pip install \
        "diffusers>=0.27.0" \
        "transformers>=4.40.0" \
        accelerate \
        protobuf \
        sentencepiece \
        soundfile \
        torchsde \
        fastapi \
        "uvicorn[standard]" \
        pydantic

WORKDIR /app
COPY server.py /app/server.py

EXPOSE 8000
CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"]