feat(agent): Optimize the entire Docker workflow and add a new Docker startup script.

Dockerfile

@@ -1,52 +1,51 @@
-# Stage 1: Base environment build
-FROM ubuntu:22.04 as builder
+# Copyright 2025 IDDM Authors
 
-# Disable interactive prompts and configure time zone
-ENV DEBIAN_FRONTEND=noninteractive \
-    TZ=Etc/UTC
-
-# Install apt-get dependencies
-RUN apt-get update && \
-    apt-get install -y --no-install-recommends ca-certificates wget git openssh-client && \
-    rm -rf /var/lib/apt/lists/*
-
-# Install Miniconda
-ARG MINICONDA_VERSION=py310_25.1.1-2
-RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-${MINICONDA_VERSION}-Linux-x86_64.sh -O ~/miniconda.sh && \
-    /bin/bash ~/miniconda.sh -b -p /opt/conda && \
-    rm ~/miniconda.sh
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
 
-# Set conda environment
-ENV PATH=/opt/conda/bin:$PATH
+#     http://www.apache.org/licenses/LICENSE-2.0
 
-# Stage 2: App deployment
-FROM ubuntu:22.04
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
-# Copy necessary content from the stage 1 (builder)
-COPY --from=builder /opt/conda /opt/conda
+# Use NVIDIA CUDA base image for GPU support
+FROM nvidia/cuda:12.8.0-cudnn-runtime-ubuntu22.04
+# For China region
+# FROM swr.cn-north-4.myhuaweicloud.com/ddn-k8s/docker.io/nvidia/cuda:12.8.0-cudnn-runtime-ubuntu22.04
 
-# Set environment variables
-ENV PATH=/opt/conda/bin:$PATH \
+# Disable interactive prompts and configure time zone
+ENV DEBIAN_FRONTEND=noninteractive \
+    TZ=Etc/UTC \
     LANG=C.UTF-8 \
     PYTHONUNBUFFERED=1
 
-# Git clone the repository
+# Install system dependencies
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends python3 python3-pip git && \
+    rm -rf /var/lib/apt/lists/* && \
+    ln -sf /usr/bin/python3 /usr/bin/python
+
+# Clone the repository
 RUN git clone --depth 1 https://github.com/chairc/Integrated-Design-Diffusion-Model.git /app
+# For China region
+# RUN git clone --depth 1 https://gitee.com/chairc/Integrated-Design-Diffusion-Model.git /app
 
 # Set working directory
 WORKDIR /app
 
-# Install (conda + pip)
-# If pip is not available, use conda
-RUN conda install -y pytorch=1.13.0 torchvision=0.14.0 torchaudio=0.13.0 pytorch-cuda=11.8 -c pytorch -c nvidia && \
-    conda clean -ya
-
-# If conda is not available, use pip
-# RUN pip install torch==1.13.0+cu116 torchvision==0.14.0+cu116 -f https://download.pytorch.org/whl/torch_stable.html
+# Install PyTorch (compatible with CUDA 12.8)
+RUN pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
+# For China region
+# RUN pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
 
-RUN cd /app/Integrated-Design-Diffusion-Model && \
-    pip install -r requirements.txt && \
-    pip cache purge
+# Install project dependencies
+RUN pip install --no-cache-dir --trusted-host pypi.org --trusted-host files.pythonhosted.org --no-deps -r requirements.txt
+# For China region
+# RUN pip install --no-cache-dir -r requirements.txt --index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
 
 # Verify installation
 RUN python -c "import torch; print(f'PyTorch version: {torch.__version__}')" && \

README.md

@@ -82,6 +82,22 @@ python train.py --xxx xxx # Replace --xxx with your training parameters
 python generate.py --xxx xxx # Replace --xxx with your generation parameters
 ```
 
+#### Docker Running
+
+Modification contents, please refer to the `docker_run.sh` script.
+
+```bash
+# Mount and build the entire project
+bash docker_run.sh -a /path/to/project
+# Mount and build specified paths
+bash docker_run.sh \
+    -d /path/to/datasets \
+    -r /path/to/results  \
+    -w /path/to/weights
+```
+
+`/path/to/project`, `/path/to/datasets`, `/path/to/results` and `/path/to/weights` are local paths, used to mount the full project, dataset, result and weight directories respectively.
+
 #### Installation
 
 In addition to running locally, there are also the following two approachs for installing this code.
@@ -96,19 +112,19 @@ The following  packages are required.
 
 ```yaml
 coloredlogs==15.0.1
-gradio>=6.0.0
+gradio==6.0.0
 matplotlib==3.7.1
-numpy>=1.25.0
-Pillow>=12.2.0
-Requests>=2.32.4
+numpy==1.25.0
+Pillow==12.2.0
+Requests==2.32.4
 scikit-image==0.22.0
 torch_summary==1.4.5
 tqdm==4.66.3
 pytorch_fid==0.3.0
 fastapi==0.115.6
 tensorboard==2.19.0
 tensorboardX==2.6.1
-transformers>=5.0.0
+transformers==5.0.0
 
 # If you want to use flash attention, please install flash-attn.
 # Compile your own environment: pip install flash-attn --no-build-isolation

README_zh.md

@@ -82,6 +82,22 @@ python train.py --xxx xxx # 用个人参数替换-xxx
 python generate.py --xxx xxx # 用个人参数替换-xxx
 ```
 
+#### Docker运行
+
+运行`构建脚本+运行容器`bash脚本，详细修改内容请参考`docker_run.sh`脚本，如遇下载问题请修改`Dockerfile`的中国区域下载地址
+
+```bash
+# 构建项目挂载
+bash docker_run.sh -a /path/to/project
+# 构建指定路径挂载
+bash docker_run.sh \
+    -d /path/to/datasets \
+    -r /path/to/results  \
+    -w /path/to/weights
+```
+
+其中，`/path/to/project`、`/path/to/datasets`、`/path/to/results`和`/path/to/weights`为本地的路径，用于挂载全项目、数据集、结果和权重目录
+
 #### 安装
 
 除了本地运行外，也可采用下列两种方式安装本代码
@@ -96,19 +112,19 @@ pip install iddm
 
 ```yaml
 coloredlogs==15.0.1
-gradio>=6.0.0
+gradio==6.0.0
 matplotlib==3.7.1
-numpy>=1.25.0
-Pillow>=12.2.0
-Requests>=2.32.4
+numpy==1.25.0
+Pillow==12.2.0
+Requests==2.32.4
 scikit-image==0.22.0
 torch_summary==1.4.5
 tqdm==4.66.3
 pytorch_fid==0.3.0
 fastapi==0.115.6
 tensorboard==2.19.0
 tensorboardX==2.6.1
-transformers>=5.0.0
+transformers==5.0.0
 
 # 如果您想使用 flash attention，请安装 flash-attn
 # 编译自己的环境：pip install flash-attn --no-build-isolation

docker_run.sh

@@ -0,0 +1,106 @@
+#!/bin/bash
+
+# Copyright 2025 IDDM Authors
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -e
+
+IMAGE_NAME="iddm-image"
+CONTAINER_NAME="iddm-container"
+
+HOST_DATASET_DIR=""
+HOST_OUTPUT_DIR=""
+HOST_WEIGHT_DIR=""
+HOST_PROJECT_DIR=""
+
+show_help() {
+    echo "Usage: bash docker_run.sh [options]"
+    echo ""
+    echo "Options:"
+    echo "  -a, --all       Host project root directory (mapped to /app)"
+    echo "  -d, --dataset   Host dataset directory      (mapped to /app/datasets)"
+    echo "  -r, --result    Host result directory        (mapped to /app/results)"
+    echo "  -w, --weight    Host weight directory        (mapped to /app/weights)"
+    echo "  -b, --build     Only build image, do not run container"
+    echo "  -h, --help      Show this help message"
+    echo ""
+    echo "Example:"
+    echo "  bash docker_run.sh -a /home/user/project"
+    echo "  bash docker_run.sh -d /home/user/datasets -r /home/user/results -w /home/user/weights"
+}
+
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        -a|--all)
+            HOST_PROJECT_DIR="$2"; shift 2 ;;
+        -d|--dataset)
+            HOST_DATASET_DIR="$2"; shift 2 ;;
+        -r|--result)
+            HOST_OUTPUT_DIR="$2"; shift 2 ;;
+        -w|--weight)
+            HOST_WEIGHT_DIR="$2"; shift 2 ;;
+        -b|--build)
+            BUILD_ONLY=true; shift ;;
+        -h|--help)
+            show_help; exit 0 ;;
+        *)
+            echo "Unknown option: $1"; show_help; exit 1 ;;
+    esac
+done
+
+echo "=============================="
+echo " Step 1: Building Docker image"
+echo "=============================="
+sudo docker build -t "${IMAGE_NAME}" .
+echo "Image '${IMAGE_NAME}' built successfully."
+echo ""
+
+if [ "${BUILD_ONLY}" = true ]; then
+    echo "Build-only mode. Exiting."
+    exit 0
+fi
+
+MOUNT_ARGS=""
+if [ -n "${HOST_PROJECT_DIR}" ]; then
+    if [ -n "${HOST_DATASET_DIR}" ] || [ -n "${HOST_OUTPUT_DIR}" ] || [ -n "${HOST_WEIGHT_DIR}" ]; then
+        echo "Warning: -a/--all is specified, individual mounts (-d/-r/-w) will be ignored."
+    fi
+    mkdir -p "${HOST_PROJECT_DIR}"
+    MOUNT_ARGS="${MOUNT_ARGS} -v ${HOST_PROJECT_DIR}:/app"
+    echo "Mount project: ${HOST_PROJECT_DIR} -> /app"
+else
+    if [ -n "${HOST_DATASET_DIR}" ]; then
+        mkdir -p "${HOST_DATASET_DIR}"
+        MOUNT_ARGS="${MOUNT_ARGS} -v ${HOST_DATASET_DIR}:/app/datasets"
+        echo "Mount dataset: ${HOST_DATASET_DIR} -> /app/datasets"
+    fi
+    if [ -n "${HOST_OUTPUT_DIR}" ]; then
+        mkdir -p "${HOST_OUTPUT_DIR}"
+        MOUNT_ARGS="${MOUNT_ARGS} -v ${HOST_OUTPUT_DIR}:/app/results"
+        echo "Mount results: ${HOST_OUTPUT_DIR} -> /app/results"
+    fi
+    if [ -n "${HOST_WEIGHT_DIR}" ]; then
+        mkdir -p "${HOST_WEIGHT_DIR}"
+        MOUNT_ARGS="${MOUNT_ARGS} -v ${HOST_WEIGHT_DIR}:/app/weights"
+        echo "Mount weights: ${HOST_WEIGHT_DIR} -> /app/weights"
+    fi
+fi
+echo ""
+
+echo "=============================="
+echo " Step 2: Running container"
+echo "=============================="
+sudo docker run --gpus all -it --rm --name "${CONTAINER_NAME}" \
+    ${MOUNT_ARGS} \
+    "${IMAGE_NAME}"

docs/en-US/01_get_started.md

@@ -7,8 +7,32 @@ Use the `git clone` or directly download the `zip` file of this repository's cod
 ```bash
 git clone https://github.com/chairc/Integrated-Design-Diffusion-Model.git
 cd Integrated-Design-Diffusion-Model
+# Run the project in a virtual environment (recommended)
+conda create -n iddm python=3.10
+pip install -r requirements.txt
+# Train the model
+cd iddm/tools
+python train.py --xxx xxx # Replace --xxx with your training parameters
+# Generate images
+python generate.py --xxx xxx # Replace --xxx with your generation parameters
 ```
 
+#### Docker Running
+
+Modification contents, please refer to the `docker_run.sh` script.
+
+```bash
+# Mount and build the entire project
+bash docker_run.sh -a /path/to/project
+# Mount and build specified paths
+bash docker_run.sh \
+    -d /path/to/datasets \
+    -r /path/to/results  \
+    -w /path/to/weights
+```
+
+`/path/to/project`, `/path/to/datasets`, `/path/to/results` and `/path/to/weights` are local paths, used to mount the full project, dataset, result and weight directories respectively.
+
 #### Installation
 
 > [!NOTE]
@@ -26,19 +50,19 @@ The following  packages are required.
 
 ```yaml
 coloredlogs==15.0.1
-gradio>=6.0.0
+gradio==6.0.0
 matplotlib==3.7.1
-numpy>=1.25.0
-Pillow>=12.2.0
-Requests>=2.32.4
+numpy==1.25.0
+Pillow==12.2.0
+Requests==2.32.4
 scikit-image==0.22.0
 torch_summary==1.4.5
 tqdm==4.66.3
 pytorch_fid==0.3.0
 fastapi==0.115.6
 tensorboard==2.19.0
 tensorboardX==2.6.1
-transformers>=5.0.0
+transformers==5.0.0
 
 # If you want to use flash attention, please install flash-attn.
 # Compile your own environment: pip install flash-attn --no-build-isolation