FROM nvidia/cuda:12.8.0-runtime-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV MODEL_DIR=/app

RUN apt-get update && apt-get install -y --no-install-recommends \
        python3.10 python3.10-dev python3-pip \
        ffmpeg gcc libc6-dev \
    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \
    && update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1 \
    && python -m pip install --upgrade pip setuptools wheel \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# Model weights
COPY high_noise_model_bnb_nf4/ /app/high_noise_model_bnb_nf4/
COPY low_noise_model_bnb_nf4/ /app/low_noise_model_bnb_nf4/
COPY models_t5_umt5-xxl-enc-bf16.pth /app/models_t5_umt5-xxl-enc-bf16.pth
COPY Wan2.1_VAE.pth /app/Wan2.1_VAE.pth
COPY tokenizer/ /app/tokenizer/

# Python packages
COPY requirements.txt /app/requirements.txt

# Pin PyTorch to 2.7.1+cu128 (must match the flash-attn prebuilt wheel)
RUN pip install --no-cache-dir \
        torch==2.7.1+cu128 torchvision torchaudio \
        --index-url https://download.pytorch.org/whl/cu128

RUN pip install --no-cache-dir -r requirements.txt

# Prebuilt flash-attn wheel: Python 3.10 / PyTorch 2.7 / CUDA 12 / cxx11abiTRUE
RUN pip install --no-cache-dir \
    https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.7cxx11abiTRUE-cp310-cp310-linux_x86_64.whl

# Application source code
COPY wan/ /app/wan/
COPY generate_prequant.py /app/generate_prequant.py
COPY load_prequant.py /app/load_prequant.py
COPY handler.py /app/handler.py

CMD ["gunicorn", "handler:app", \
    "--bind", "0.0.0.0:8080", \
    "--workers", "1", \
    "--threads", "4", \
    "--timeout", "0", \
    "--access-logfile", "-", \
    "--error-logfile", "-"]