# 基础镜像：CUDA 12.4 (匹配 RTX 50 系)
FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04

# 环境变量
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
ENV DEBIAN_FRONTEND=noninteractive
ENV HF_HOME=/root/.cache/huggingface

WORKDIR /app

# 1. 系统依赖
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3-pip \
    git \
    wget \
    vim \
    libgl1 \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

# 2. 设置 pip 清华源全局配置 (省去每次敲参数)
RUN pip3 config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

# 3. 🔥 核心：RTX 5060 专用 PyTorch (cu128)
# 必须显式指定 index-url 覆盖上面的清华源配置，因为 cu128 只有官方 nightly 有
RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

# 4. 安装 DiffSim/Diffusion 核心库
# 包含 tqdm 用于显示进度条
RUN pip install \
    "diffusers>=0.24.0" \
    "transformers>=4.35.0" \
    accelerate \
    scipy \
    safetensors \
    opencv-python \
    ftfy \
    regex \
    timm \
    einops \
    lpips \
    tqdm \
    matplotlib

# 5. 默认指令：启动 bash，让你进入终端
CMD ["/bin/bash"]