Dockerfile.gpu

# Woosh-DFlow text-to-audio BACKEND Space -- GPU build (dedicated-GPU Space).
#
# To use this on a GPU Space, rename this file to `Dockerfile`. It is identical
# to the CPU Dockerfile except for the torch wheel index (CUDA cu128, matching
# Woosh's pyproject [tool.uv.index] "pytorch-cu128"). HF GPU Spaces provide the
# NVIDIA driver, so the CUDA torch wheels run without a CUDA base image.
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 -- CUDA 12.8 wheels (this is the GPU image).
RUN pip install --upgrade pip \
 && pip install --index-url https://download.pytorch.org/whl/cu128 \
        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.1" \
        "hydra-core>=1.3.2" \
        "lightning>=2.5.6" \
        "timm>=1.0.22" \
        "torchdiffeq>=0.2.5" \
        "transformers>=4.57.2" \
        "pydantic>=2.12.4" \
        "omegaconf>=2.3.0" \
        "av>=16.1.0" \
        "requests" \
        "soundfile>=0.13.1" \
        "gradio>=6.9.0" \
 && pip install --no-deps "hear21passt==0.0.26" \
 && pip install --no-deps "git+https://github.com/SonyResearch/Woosh.git@main"

# 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).
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://github.com/SonyResearch/Woosh/releases/download/v1.0.0/$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"]