[SD] v4.0 cleanup and bug fixes

mlcommons · Feb 12, 2024 · a97d3aa · a97d3aa
1 parent 00f04c5
commit a97d3aa
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 8 deletions.
diff --git a/stable_diffusion/Dockerfile b/stable_diffusion/Dockerfile
@@ -1,4 +1,4 @@
-ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:22.12-py3
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:24.01-py3
 FROM ${FROM_IMAGE_NAME}
 
 ENV DEBIAN_FRONTEND=noninteractive
@@ -7,6 +7,11 @@ ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update
 RUN apt-get install -y ffmpeg libsm6 libxext6
 
+# Remove the opencv version shipped with the base image
+# https://github.com/opencv/opencv-python/issues/884
+RUN pip uninstall -y opencv
+RUN rm -rf /usr/local/lib/python3.10/dist-packages/cv2/
+
 # install LDM
 COPY . /diffusion
 RUN cd /diffusion && \

diff --git a/stable_diffusion/requirements.txt b/stable_diffusion/requirements.txt
@@ -1,7 +1,7 @@
 albumentations==1.3.0
-opencv-python==4.7.0.72
+opencv-python-headless==4.9.0.80
 pudb==2019.2
-prefetch_generator
+prefetch_generator==1.0.3
 imageio==2.9.0
 imageio-ffmpeg==0.4.2
 torchmetrics==0.11.4
@@ -20,6 +20,6 @@ colossalai==0.2.7
 invisible-watermark==0.1.5
 diffusers==0.14.0
 cloudpathlib==0.13.0
-git+https://github.com/facebookresearch/xformers.git@5eb0dbf315d14b5f7b38ac2ff3d8379beca7df9b#egg=xformers
+xformers==0.0.24
 bitsandbytes==0.37.2
-git+https://github.com/mlcommons/logging.git@8405a08bbfc724f8888c419461c02d55a6ac960c
+git+https://github.com/ahmadki/logging.git@6328dfcd04a1bda8ae5b3e4f1f586e9037d42ddf
diff --git a/stable_diffusion/run_and_time.sh b/stable_diffusion/run_and_time.sh
@@ -4,7 +4,7 @@
 : "${GPUS_PER_NODE:=8}"
 : "${CHECKPOINT:=/checkpoints/sd/512-base-ema.ckpt}"
 : "${RESULTS_DIR:=}"
-: "${CONFIG:=./configs/train_512_latents.yaml}"
+: "${CONFIG:=./configs/train_01x08x08.yaml}"
 
 while [ "$1" != "" ]; do
     case $1 in

diff --git a/stable_diffusion/scripts/docker/build.sh b/stable_diffusion/scripts/docker/build.sh
@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 
-: "${SRC_IMG:=nvcr.io/nvidia/pytorch:22.12-py3}"
-: "${DST_IMG:=mlperf_sd:22.12-py3}"
+: "${SRC_IMG:=nvcr.io/nvidia/pytorch:24.01-py3}"
+: "${DST_IMG:=mlperf_sd:24.01-py3}"
 
 while [ "$1" != "" ]; do
     case $1 in