build: fix torch build to work with the installation of vision models

Author: fracape

.gitlab-ci.yml

@@ -47,7 +47,7 @@ ruff-lint:
   before_script:
     - *default-before-script
     - uv pip install -U pip wheel setuptools
-    - uv sync --group=dev
+    - uv sync --group=dev --extra=cpu
   script:
     - uv run --no-sync make check-ruff-lint
   tags:
@@ -64,7 +64,7 @@ pytest-py38:
     - *default-before-script
     - uv pip install -U pip wheel setuptools
     - bash scripts/install_uv.sh --no-weights --cpu -m "detectron2,yolox,segment_anything,jde"
-    - uv sync --inexact --extra=dev
+    - uv sync --inexact --extra=dev --extra=cpu
   script:
     - uv run --no-sync pytest
   tags:
@@ -81,7 +81,7 @@ pytest-py310:
     - *default-before-script
     - uv pip install -U pip wheel setuptools
     - bash scripts/install_uv.sh --no-weights --cpu -m "detectron2,yolox,segment_anything,jde"
-    - uv sync --inexact --extra=dev
+    - uv sync --inexact --extra=dev --extra=cpu
   script:
     - uv run --no-sync pytest
   tags:
@@ -96,7 +96,7 @@ pytest-pip-py38:
     - python3 -m venv venv
     - source venv/bin/activate
     - pip install -U pip wheel setuptools
-    - pip install torch==2.0.0 torchvision==0.15.1 --index-url https://download.pytorch.org/whl/cpu
+    - pip install torch==2.0.0 torchvision==0.15.1 --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
     - bash scripts/install.sh --no-weights --cpu -m "detectron2,yolox,segment_anything,jde"
     - pip install -e .[dev,cpu]
   script:
@@ -126,7 +126,7 @@ ruff-format:
   before_script:
     - *default-before-script
     - uv pip install -U pip wheel setuptools
-    - uv sync --group=dev
+    - uv sync --group=dev --extra=cpu
   script:
     - uv run --no-sync make check-ruff-format
   tags:
@@ -138,7 +138,7 @@ ruff-organize-imports:
   before_script:
     - *default-before-script
     - uv pip install -U pip wheel setuptools
-    - uv sync --group=dev
+    - uv sync --group=dev --extra=cpu
   script:
     - uv run --no-sync make check-ruff-organize-imports
   tags:
@@ -150,7 +150,7 @@ doc:
   before_script:
     - *default-before-script
     - uv pip install -U pip wheel setuptools
-    - uv sync --group=doc
+    - uv sync --group=doc --extra=cpu
     - cd docs
   script:
     - uv run --no-sync make html

scripts/install.sh

@@ -258,7 +258,9 @@ detect_cuda_version () {
 }
 
 install_torch () {
-    "${PIP[@]}" install "torch==${TORCH_VERSION}" "torchvision==${TORCHVISION_VERSION}" --index-url "https://download.pytorch.org/whl/${BUILD_SUFFIX}"
+    "${PIP[@]}" install "torch==${TORCH_VERSION}" "torchvision==${TORCHVISION_VERSION}" \
+        --index-url "https://download.pytorch.org/whl/${BUILD_SUFFIX}" \
+        --extra-index-url "https://pypi.org/simple"
 
 }
 
@@ -294,7 +296,7 @@ install_detectron2 () {
         "${PIP[@]}" install --no-build-isolation -e .
     elif [[ "${PACKAGE_MANAGER}" == "uv" ]]; then
         cd "${COMPRESSAI_VISION_ROOT_DIR}"
-        uv sync --inexact --group=models-detectron2
+        uv sync --inexact --group=models-detectron2 --extra="${BUILD_SUFFIX}"
     fi
 
     cd "${COMPRESSAI_VISION_ROOT_DIR}"
@@ -367,7 +369,7 @@ install_jde () {
         "${PIP[@]}" install --no-build-isolation -e .
     elif [[ "${PACKAGE_MANAGER}" == "uv" ]]; then
         cd "${COMPRESSAI_VISION_ROOT_DIR}"
-        uv sync --inexact --group=models-jde
+        uv sync --inexact --group=models-jde --extra="${BUILD_SUFFIX}"
     fi
 
     cd "${COMPRESSAI_VISION_ROOT_DIR}"
@@ -403,7 +405,7 @@ install_yolox () {
         "${PIP[@]}" install --no-build-isolation -e .
     elif [[ "${PACKAGE_MANAGER}" == "uv" ]]; then
         cd "${COMPRESSAI_VISION_ROOT_DIR}"
-        uv sync --inexact --group=models-yolox
+        uv sync --inexact --group=models-yolox --extra="${BUILD_SUFFIX}"
     fi
 
     cd "${COMPRESSAI_VISION_ROOT_DIR}"
@@ -442,7 +444,7 @@ install_mmpose () {
         "${PIP[@]}" install isort==5.13.2
     elif [[ "${PACKAGE_MANAGER}" == "uv" ]]; then
         cd "${COMPRESSAI_VISION_ROOT_DIR}"
-        uv sync --inexact --group=models-mmpose
+        uv sync --inexact --group=models-mmpose --extra="${BUILD_SUFFIX}"
     fi
 
     "${MIM[@]}" install "mmcv==2.0.1"
@@ -478,7 +480,7 @@ install_segment_anything () {
         "${PIP[@]}" install -e .
     elif [[ "${PACKAGE_MANAGER}" == "uv" ]]; then
         cd "${COMPRESSAI_VISION_ROOT_DIR}"
-        uv sync --inexact --group=models-segment-anything
+        uv sync --inexact --group=models-segment-anything --extra="${BUILD_SUFFIX}"
     fi
 
     cd "${COMPRESSAI_VISION_ROOT_DIR}"
@@ -513,7 +515,7 @@ install_sam2 () {
         "${PIP[@]}" install -e .
     elif [[ "${PACKAGE_MANAGER}" == "uv" ]]; then
         cd "${COMPRESSAI_VISION_ROOT_DIR}"
-        uv sync --inexact --group=models-sam2
+        uv sync --inexact --group=models-sam2 --extra="${BUILD_SUFFIX}"
     fi
 
     cd "${COMPRESSAI_VISION_ROOT_DIR}"
@@ -568,4 +570,4 @@ download_weights () {
 
 if [[ "${__SOURCE_ONLY__:-0}" -eq 0 ]]; then
     main "$@"
-fi
+fi