Update Dockerfile

Dockerfile

@@ -1,5 +1,6 @@
 # ==== GPU base (CUDA 12.1 + cuDNN) ====
-FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+# Use *devel* so nvcc is available for packages like DeepSpeed/xFormers to compile
+FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONDONTWRITEBYTECODE=1 \
@@ -9,33 +10,33 @@ ENV DEBIAN_FRONTEND=noninteractive \
     TRANSFORMERS_CACHE=/root/.cache/huggingface/transformers \
     GRADIO_SERVER_NAME=0.0.0.0 \
     GRADIO_SERVER_PORT=8080 \
-    DISPLAY=:99
+    DISPLAY=:99 \
+    CUDA_HOME=/usr/local/cuda
 
-# ---- System deps (ffmpeg + headless GL/X) ----
+# ---- System deps (ffmpeg + headless GL/X) + build toolchain for CUDA/JIT ----
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    python3 python3-pip python3-venv \
+    python3 python3-pip python3-venv python3-dev \
     git git-lfs \
     ffmpeg \
-    libgl1 \
-    libgl1-mesa-dri \
-    libglib2.0-0 \
-    xvfb \
+    libgl1 libgl1-mesa-dri libglib2.0-0 xvfb \
+    build-essential ninja-build cmake libaio-dev \
  && git lfs install \
  && rm -rf /var/lib/apt/lists/*
 
-# ---- Workdir & copy requirements first (cache-friendly) ----
 WORKDIR /app
 COPY requirements.txt /app/requirements.txt
 
-# ---- Python deps (Torch CUDA 12.1 first) ----
+# ---- Python deps (install PyTorch 2.x with cu121 first) ----
+# Note: we install 'ninja' (Python) too because PyTorch/DS use it for JIT builds.
 RUN python3 -m pip install --upgrade pip \
  && python3 -m pip install --extra-index-url https://download.pytorch.org/whl/cu121 \
       torch torchvision torchaudio \
- && python3 -m pip install -r requirements.txt \
+ && python3 -m pip install ninja \
+ && DS_BUILD_AIO=0 DS_BUILD_SPARSE_ATTN=0 TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0" \
+      python3 -m pip install -r requirements.txt \
  && python3 -m pip install huggingface_hub imageio-ffmpeg
 
 # ---- App code ----
-# Expect: gradio_app.py (updated), image_to_video.py, demo/, etc.
 COPY . /app
 
 # ---- Ports & default envs ----
@@ -49,12 +50,5 @@ ENV DEVICE=cuda \
     DEPTH_MODEL_PATH=pretrained/Metric3D/metric_depth_vit_large_800k.pth \
     CAPTION_MODEL_PATH=pretrained/Qwen2.5-VL-7B-Instruct
 
-# Optional (set at runtime to auto-download via gradio_app.py):
-# ENV REPO_COGVIDEOX=THUDM/CogVideoX1.5-5B-I2V
-# ENV REPO_METRIC3D=JUGGHM/Metric3D
-# ENV REPO_QWEN_VL=Qwen/Qwen2.5-VL-7B-Instruct
-# ENV REPO_REALCAM=MuteApo/RealCam-I2V
-# ENV HF_TOKEN=   # pass at runtime; don't bake secrets
-
 # ---- Launch under Xvfb to satisfy headless GL ----
 CMD ["bash", "-lc", "xvfb-run -s '-screen 0 1280x720x24' python3 finetune/gradio_app.py"]