# 1. 基础镜像：使用你验证过的 CUDA 12.4 开发版
FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04

# 环境变量
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
ENV DEBIAN_FRONTEND=noninteractive
# 针对 RTX 50 系列的架构优化 (Blackwell = sm_100/sm_120，这里涵盖所有新架构)
ENV TORCH_CUDA_ARCH_LIST="9.0;10.0+PTX"

# 2. 安装系统工具 (增加了 ffmpeg 用于视频处理)
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3-pip \
    git \
    wget \
    ffmpeg \
    libgl1 \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    build-essential \
    ninja-build \
    && rm -rf /var/lib/apt/lists/*

# 建立 python 软链接，方便直接用 python 命令
RUN ln -s /usr/bin/python3.10 /usr/bin/python

WORKDIR /app

# 3. 升级 pip (使用清华源加速)
RUN python3 -m pip install --upgrade pip -i https://pypi.tuna.tsinghua.edu.cn/simple --default-timeout=100

# 4. 🔥 核心修复：安装适配 RTX 50 系列的 PyTorch (cu128)
# 这是解决 "sm_120 is not compatible" 的关键一步
RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 --default-timeout=100

# 5. 安装 SAM3 依赖 (包含之前报错缺少的 decord, pycocotools)
# 移除了 transformers 版本限制，使用最新版适配 SAM3
RUN pip install \
    opencv-python-headless \
    matplotlib \
    jupyter \
    jupyterlab \
    ipympl \
    pyyaml \
    tqdm \
    hydra-core \
    iopath \
    pillow \
    networkx \
    scipy \
    pandas \
    timm \
    einops \
    transformers \
    tokenizers \
    decord \
    pycocotools \
    -i https://pypi.tuna.tsinghua.edu.cn/simple \
    --default-timeout=100

# 6. 拉取 SAM3 代码
RUN git clone https://github.com/facebookresearch/sam3.git sam3_code

# 7. 安装 SAM3 包
WORKDIR /app/sam3_code
RUN pip install -e .

# 8. 设置工作目录和入口
WORKDIR /app
EXPOSE 8888

# 默认保持运行，方便进入终端调试
CMD ["/bin/bash", "-c", "jupyter lab --ip=0.0.0.0 --port=8888 --allow-root --no-browser --NotebookApp.token='sam3'"]