FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    UV_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu121 \
    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128

# Base OS tools + media stack + Python toolchain.
RUN apt-get update && apt-get install -y --no-install-recommends \
    ffmpeg \
    git \
    git-lfs \
    ca-certificates \
    curl \
    build-essential \
    pkg-config \
    ninja-build \
    libglib2.0-0 \
    libgl1 \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sf /usr/bin/python3.10 /usr/bin/python \
    && git lfs install

# Install uv.
COPY --from=ghcr.io/astral-sh/uv:0.6.17 /uv /uvx /usr/local/bin/
RUN uv python install 3.10
# Place executables in the environment at the front of the path
ENV PATH="/app/.venv/bin:$PATH"

WORKDIR /app

# Install app Python dependencies first for better layer caching.
COPY pyproject.toml README.md /app/

RUN uv sync --no-dev --no-install-project

# Ensure HunyuanVideo source exists in the image.
ARG HUNYUAN_REPO=https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5.git
RUN if [ ! -f /app/HunyuanVideo-1.5/requirements.txt ]; then \
      rm -rf /app/HunyuanVideo-1.5 && \
      git clone --depth 1 "$HUNYUAN_REPO" /app/HunyuanVideo-1.5; \
    fi

# Install HunyuanVideo dependencies from upstream README guidance.
RUN uv pip install --index-strategy unsafe-best-match -r /app/HunyuanVideo-1.5/requirements.txt \
    && uv pip install --upgrade tencentcloud-sdk-python \
    && uv pip install sgl-kernel==0.3.18

# Optional attention backends from Hunyuan docs.
# Build with: --build-arg INSTALL_OPTIONAL_ATTENTION=1
ARG INSTALL_OPTIONAL_ATTENTION=0
RUN if [ "$INSTALL_OPTIONAL_ATTENTION" = "1" ]; then \
    uv pip install flash-attn --no-build-isolation && \
      git clone --depth 1 https://github.com/Tencent-Hunyuan/flex-block-attn.git /tmp/flex-block-attn && \
      cd /tmp/flex-block-attn && git submodule update --init --recursive && python setup.py install && \
      git clone --depth 1 https://github.com/cooper1637/SageAttention.git /tmp/SageAttention && \
      cd /tmp/SageAttention && python setup.py install; \
    fi

# Copy application source after dependencies are installed.
COPY . /app

# Default pipeline entrypoint.
CMD ["uv", "run", "run_video_pipeline.py"]