# Woosh-DFlow text-to-audio Space for HARP -- CPU build (free HF Space).
#
# Woosh and PyHARP run in one process here. They were split across two Spaces
# while PyHARP pinned Gradio 5.28, which Woosh's Gradio 6 requirement could not
# meet; PyHARP v0.4.1 requires Gradio 6 itself, so the split is no longer
# needed and the HARP endpoint is served from the same image as the model.
#
# For a dedicated-GPU Space, swap the torch wheel index below from cpu to the
# matching CUDA one. ZeroGPU is not an option: it is offered to Gradio SDK
# Spaces only, and the weights have to be baked in rather than downloaded on
# each cold start, which is what makes this a Docker Space.
#
# The install sequence mirrors Woosh's pyproject.toml, its lower bounds pinned to the
# versions they resolved to, with two upstream quirks:
#   * hear21passt is installed with --no-deps (Woosh's uv config overrides its
#     requires-dist to empty; its runtime imports are covered by the deps below).
#   * The woosh package itself is installed with --no-deps (its deps are pinned
#     explicitly here) so pip never re-resolves the torch stack.
FROM python:3.12-slim-bookworm

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    HF_HOME=/tmp/hf \
    PORT=7860

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential git curl unzip ffmpeg libsndfile1 \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# 1) Torch stack -- CPU wheels (this is the CPU image).
RUN pip install --upgrade pip \
 && pip install --index-url https://download.pytorch.org/whl/cpu \
        torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0

# 2) Woosh's declared runtime deps (from pyproject.toml), then Woosh itself.
RUN pip install \
        "einops==0.8.2" \
        "hydra-core==1.3.7" \
        "lightning==2.6.6" \
        "timm==1.0.30" \
        "torchdiffeq==0.2.5" \
        "transformers==5.20.2" \
        "pydantic==2.14.1" \
        "omegaconf==2.3.1" \
        "av==19.0.2" \
        "requests==2.34.2" \
        "soundfile==0.14.0" \
        "gradio==6.24.0" \
 && pip install --no-deps "hear21passt==0.0.26" \
 && pip install --no-deps "git+https://github.com/SonyResearch/Woosh.git@7f238bbd5d6dbc94d27190c07827b2d0eaefa268"

# 2b) PyHARP, which serves the HARP endpoint from this same process. It pulls
#     Gradio 6 and descript-audiotools, whose torch requirement the CPU wheels
#     above already satisfy, so pip leaves the torch stack alone. The CPU index
#     stays on hand in case a resolution ever does reach for torch.
RUN pip install --extra-index-url https://download.pytorch.org/whl/cpu \
        "git+https://github.com/TEAMuP-dev/pyharp.git@v0.4.2" "gradio==6.24.0"

# 3) Model weights (CC-BY-NC, official v1.0.0 release), baked in for fast
#    restarts. Woosh-DFlow needs the DFlow LDM + the audio autoencoder (AE) +
#    the audio-CLAP text conditioner (TextConditionerA). The zips are mirrored
#    unchanged from the GitHub release in teamup-tech/Woosh-weights, read at a
#    pinned revision. Each extracts to the correct checkpoints/<NAME> folder
#    under /app.
RUN set -eux; cd /app; \
    for A in Woosh-DFlow Woosh-AE TextConditionerA; do \
      echo "Downloading $A.zip ..."; \
      curl -fL --retry 3 -o "$A.zip" \
        "https://huggingface.co/teamup-tech/Woosh-weights/resolve/b34c6218f40b0d25bac40691f7e34b27f0246ce7/$A.zip"; \
      unzip -q "$A.zip" -d /app; \
      rm -f "$A.zip"; \
    done; \
    ls -R /app/checkpoints | head -n 60

COPY app.py /app/app.py

EXPOSE 7860
CMD ["python", "app.py"]