mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
55 lines
2.2 KiB
Bash
55 lines
2.2 KiB
Bash
#!/usr/bin/env bash
|
|
|
|
# This file is used to set environment variables for the build.sh and run.sh scripts.
|
|
|
|
# Try to detect the container flavor if no PIP_EXTRA_INDEX_URL got specified
|
|
if [[ -z "$PIP_EXTRA_INDEX_URL" ]]; then
|
|
|
|
# Activate virtual environment if not already activated and exists
|
|
if [[ -z $VIRTUAL_ENV ]]; then
|
|
[[ -e "$(dirname "${BASH_SOURCE[0]}")/../.venv/bin/activate" ]] \
|
|
&& source "$(dirname "${BASH_SOURCE[0]}")/../.venv/bin/activate" \
|
|
&& echo "Activated virtual environment: $VIRTUAL_ENV"
|
|
fi
|
|
|
|
# Decide which container flavor to build if not specified
|
|
if [[ -z "$CONTAINER_FLAVOR" ]] && python -c "import torch" &>/dev/null; then
|
|
# Check for CUDA and ROCm
|
|
CUDA_AVAILABLE=$(python -c "import torch;print(torch.cuda.is_available())")
|
|
ROCM_AVAILABLE=$(python -c "import torch;print(torch.version.hip is not None)")
|
|
if [[ "${CUDA_AVAILABLE}" == "True" ]]; then
|
|
CONTAINER_FLAVOR="cuda"
|
|
elif [[ "${ROCM_AVAILABLE}" == "True" ]]; then
|
|
CONTAINER_FLAVOR="rocm"
|
|
else
|
|
CONTAINER_FLAVOR="cpu"
|
|
fi
|
|
fi
|
|
|
|
# Set PIP_EXTRA_INDEX_URL based on container flavor
|
|
if [[ "$CONTAINER_FLAVOR" == "rocm" ]]; then
|
|
PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/rocm"
|
|
elif [[ "$CONTAINER_FLAVOR" == "cpu" ]]; then
|
|
PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
|
|
# elif [[ -z "$CONTAINER_FLAVOR" || "$CONTAINER_FLAVOR" == "cuda" ]]; then
|
|
# PIP_PACKAGE=${PIP_PACKAGE-".[xformers]"}
|
|
fi
|
|
fi
|
|
|
|
# Variables shared by build.sh and run.sh
|
|
REPOSITORY_NAME="${REPOSITORY_NAME-$(basename "$(git rev-parse --show-toplevel)")}"
|
|
REPOSITORY_NAME="${REPOSITORY_NAME,,}"
|
|
VOLUMENAME="${VOLUMENAME-"${REPOSITORY_NAME}_data"}"
|
|
ARCH="${ARCH-$(uname -m)}"
|
|
PLATFORM="${PLATFORM-linux/${ARCH}}"
|
|
INVOKEAI_BRANCH="${INVOKEAI_BRANCH-$(git branch --show)}"
|
|
CONTAINER_REGISTRY="${CONTAINER_REGISTRY-"ghcr.io"}"
|
|
CONTAINER_REPOSITORY="${CONTAINER_REPOSITORY-"$(whoami)/${REPOSITORY_NAME}"}"
|
|
CONTAINER_FLAVOR="${CONTAINER_FLAVOR-cuda}"
|
|
CONTAINER_TAG="${CONTAINER_TAG-"${INVOKEAI_BRANCH##*/}-${CONTAINER_FLAVOR}"}"
|
|
CONTAINER_IMAGE="${CONTAINER_REGISTRY}/${CONTAINER_REPOSITORY}:${CONTAINER_TAG}"
|
|
CONTAINER_IMAGE="${CONTAINER_IMAGE,,}"
|
|
|
|
# enable docker buildkit
|
|
export DOCKER_BUILDKIT=1
|