Revert "Remove vision packages from CI scripts (#125546)"

This reverts commit f42ea14c3f795082138421fcef90d24f64c6fd35.

Reverted https://github.com/pytorch/pytorch/pull/125546 on behalf of https://github.com/huydhn due to I think we are using vision in inductor tests with their various models there ([comment](https://github.com/pytorch/pytorch/pull/125546#issuecomment-2105174723))

.ci/docker/build.sh

@@ -98,6 +98,7 @@ case "$image" in
     GCC_VERSION=9
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
     UCC_COMMIT=${_UCC_COMMIT}
@@ -111,6 +112,7 @@ case "$image" in
     GCC_VERSION=9
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
     UCC_COMMIT=${_UCC_COMMIT}
@@ -125,6 +127,7 @@ case "$image" in
     GCC_VERSION=9
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
     UCC_COMMIT=${_UCC_COMMIT}
@@ -138,6 +141,7 @@ case "$image" in
     GCC_VERSION=9
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
     UCC_COMMIT=${_UCC_COMMIT}
@@ -149,6 +153,7 @@ case "$image" in
     CLANG_VERSION=10
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     CONDA_CMAKE=yes
     ONNX=yes
     ;;
@@ -167,6 +172,7 @@ case "$image" in
     CLANG_VERSION=10
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     VULKAN_SDK_VERSION=1.2.162.1
     SWIFTSHADER=yes
     CONDA_CMAKE=yes
@@ -177,6 +183,7 @@ case "$image" in
     CLANG_VERSION=10
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     VULKAN_SDK_VERSION=1.2.162.1
     SWIFTSHADER=yes
     CONDA_CMAKE=yes
@@ -187,6 +194,7 @@ case "$image" in
     GCC_VERSION=9
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     CONDA_CMAKE=yes
     TRITON=yes
     ;;
@@ -195,6 +203,7 @@ case "$image" in
     GCC_VERSION=9
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     ROCM_VERSION=6.0
     NINJA_VERSION=1.9.0
     CONDA_CMAKE=yes
@@ -205,6 +214,7 @@ case "$image" in
     GCC_VERSION=9
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     ROCM_VERSION=6.1
     NINJA_VERSION=1.9.0
     CONDA_CMAKE=yes
@@ -215,6 +225,7 @@ case "$image" in
     GCC_VERSION=11
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     BASEKIT_VERSION=2024.0.0-49522
     NINJA_VERSION=1.9.0
     CONDA_CMAKE=yes
@@ -225,6 +236,7 @@ case "$image" in
     GCC_VERSION=11
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     KATEX=yes
     CONDA_CMAKE=yes
     TRITON=yes
@@ -238,6 +250,7 @@ case "$image" in
     CLANG_VERSION=12
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     TRITON=yes
     ;;
   pytorch-linux-jammy-py3-clang12-asan)
@@ -245,6 +258,7 @@ case "$image" in
     CLANG_VERSION=12
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     CONDA_CMAKE=yes
     TRITON=yes
     ;;
@@ -252,12 +266,14 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.10
     CLANG_VERSION=15
     CONDA_CMAKE=yes
+    VISION=yes
     ;;
   pytorch-linux-jammy-py3.8-gcc11)
     ANACONDA_PYTHON_VERSION=3.8
     GCC_VERSION=11
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     KATEX=yes
     CONDA_CMAKE=yes
     TRITON=yes
@@ -288,6 +304,7 @@ case "$image" in
     ACL=yes
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     CONDA_CMAKE=yes
     # snadampal: skipping sccache due to the following issue
     # https://github.com/pytorch/pytorch/issues/121559
@@ -300,6 +317,7 @@ case "$image" in
     # Catch-all for builds that are not hardcoded.
     PROTOBUF=yes
     DB=yes
+    VISION=yes
     echo "image '$image' did not match an existing build configuration"
     if [[ "$image" == *py* ]]; then
       extract_version_from_image_name py ANACONDA_PYTHON_VERSION

.ci/docker/centos-rocm/Dockerfile

@@ -62,6 +62,13 @@ RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
 RUN rm install_db.sh
 ENV INSTALLED_DB ${DB}
 
+# (optional) Install vision packages like OpenCV and ffmpeg
+ARG VISION
+COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
+RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
+RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
+ENV INSTALLED_VISION ${VISION}
+
 # Install rocm
 ARG ROCM_VERSION
 COPY ./common/install_rocm.sh install_rocm.sh

.ci/docker/common/install_base.sh

@@ -116,6 +116,7 @@ install_centos() {
     libstdc++-devel \
     libsndfile-devel \
     make \
+    opencv-devel \
     sudo \
     wget \
     vim \

.ci/docker/common/install_vision.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+
+set -ex
+
+install_ubuntu() {
+  apt-get update
+  apt-get install -y --no-install-recommends \
+          libopencv-dev \
+          libavcodec-dev
+
+  # Cleanup
+  apt-get autoclean && apt-get clean
+  rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+}
+
+install_centos() {
+  # Need EPEL for many packages we depend on.
+  # See http://fedoraproject.org/wiki/EPEL
+  yum --enablerepo=extras install -y epel-release
+
+  yum install -y \
+      opencv-devel \
+      ffmpeg-devel
+
+  # Cleanup
+  yum clean all
+  rm -rf /var/cache/yum
+  rm -rf /var/lib/yum/yumdb
+  rm -rf /var/lib/yum/history
+}
+
+# Install base packages depending on the base OS
+ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
+case "$ID" in
+  ubuntu)
+    install_ubuntu
+    ;;
+  centos)
+    install_centos
+    ;;
+  *)
+    echo "Unable to determine OS..."
+    exit 1
+    ;;
+esac
+
+# Cache vision models used by the test
+source "$(dirname "${BASH_SOURCE[0]}")/cache_vision_models.sh"

.ci/docker/ubuntu-cuda/Dockerfile

@@ -56,6 +56,13 @@ RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
 RUN rm install_db.sh
 ENV INSTALLED_DB ${DB}
 
+# (optional) Install vision packages like OpenCV and ffmpeg
+ARG VISION
+COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
+RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
+RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
+ENV INSTALLED_VISION ${VISION}
+
 # (optional) Install UCC
 ARG UCX_COMMIT
 ARG UCC_COMMIT

.ci/docker/ubuntu-rocm/Dockerfile

@@ -53,6 +53,13 @@ RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
 RUN rm install_db.sh
 ENV INSTALLED_DB ${DB}
 
+# (optional) Install vision packages like OpenCV and ffmpeg
+ARG VISION
+COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
+RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
+RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
+ENV INSTALLED_VISION ${VISION}
+
 # Install rocm
 ARG ROCM_VERSION
 COPY ./common/install_rocm.sh install_rocm.sh

.ci/docker/ubuntu-xpu/Dockerfile

@@ -83,6 +83,13 @@ RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
 RUN rm install_db.sh
 ENV INSTALLED_DB ${DB}
 
+# (optional) Install vision packages like OpenCV and ffmpeg
+ARG VISION
+COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
+RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
+RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
+ENV INSTALLED_VISION ${VISION}
+
 # (optional) Install non-default CMake version
 ARG CMAKE_VERSION
 COPY ./common/install_cmake.sh install_cmake.sh

.ci/docker/ubuntu/Dockerfile

@@ -73,6 +73,20 @@ RUN if [ -n "${PROTOBUF}" ]; then bash ./install_protobuf.sh; fi
 RUN rm install_protobuf.sh
 ENV INSTALLED_PROTOBUF ${PROTOBUF}
 
+# (optional) Install database packages like LMDB and LevelDB
+ARG DB
+COPY ./common/install_db.sh install_db.sh
+RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
+RUN rm install_db.sh
+ENV INSTALLED_DB ${DB}
+
+# (optional) Install vision packages like OpenCV and ffmpeg
+ARG VISION
+COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
+RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
+RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
+ENV INSTALLED_VISION ${VISION}
+
 # (optional) Install Android NDK
 ARG ANDROID
 ARG ANDROID_NDK

.ci/pytorch/common_utils.sh

@@ -196,7 +196,6 @@ function checkout_install_torchdeploy() {
   pushd multipy
   git checkout "${commit}"
   python multipy/runtime/example/generate_examples.py
-  sudo apt-get update && sudo apt-get install -y liblzma-dev
   BUILD_CUDA_TESTS=1 pip install -e .
   popd
   popd

.ci/pytorch/test.sh

@@ -1233,6 +1233,8 @@ elif [[ "${TEST_CONFIG}" == *torchbench* ]]; then
   install_torchtext
   install_torchvision
   id=$((SHARD_NUMBER-1))
+  # https://github.com/opencv/opencv-python/issues/885
+  pip_install opencv-python==4.8.0.74
   if [[ "${TEST_CONFIG}" == *inductor_torchbench_smoketest_perf* ]]; then
     checkout_install_torchbench hf_Bert hf_Albert nanogpt timm_vision_transformer
     PYTHONPATH=$(pwd)/torchbench test_inductor_torchbench_smoketest_perf