FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV HF_HOME=/root/.cache/huggingface
ENV HF_HUB_DISABLE_TELEMETRY=1

# system deps
RUN apt-get update && apt-get install -y --no-install-recommends \
        git curl ffmpeg libgl1 libglib2.0-0 && \
    rm -rf /var/lib/apt/lists/*

# Wan2.1 official inference code
RUN git clone --depth 1 https://github.com/Wan-Video/Wan2.1.git /opt/Wan2.1

# python deps (torch 2.5.1+cu124 is preinstalled in the base image)
RUN pip install --no-cache-dir \
        "xfuser==0.4.5" \
        "yunchang>=0.6.0" \
        "diffusers>=0.33.0" \
        "transformers>=4.49.0" \
        "tokenizers>=0.20.3" \
        "accelerate>=1.1.1" \
        "tqdm" \
        "imageio" \
        "easydict" \
        "ftfy" \
        "dashscope" \
        "imageio-ffmpeg" \
        "opencv-python>=4.9.0.80" \
        "numpy>=1.23.5,<2" \
        "gradio>=5.0.0" \
        "huggingface_hub[cli]"

# flash-attn prebuilt wheel (required by the xfuser USP attention backend)
RUN pip install --no-cache-dir \
    https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3.post1/flash_attn-2.8.3.post1+cu12torch2.5cxx11abiFALSE-cp311-cp311-linux_x86_64.whl

# Wan-AI/Wan2.1-T2V-1.3B checkpoints (~17.5GB) baked into the image
RUN huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B --local-dir /opt/Wan2.1-T2V-1.3B

COPY app.py /opt/app/app.py

WORKDIR /opt/Wan2.1

ENV NCCL_DEBUG=INFO

ENTRYPOINT ["torchrun", "--nproc_per_node=8", "--master_port=7861", "/opt/app/app.py"]