tdpeuter/open-webui
tdpeuter/open-webui
1
Fork 0
forked from open-webui/open-webui
Code Issues Pull requests 1 Projects 1 Releases 2 Packages 1 Activity Actions

cuda fix

Browse source
This commit is contained in:
Jannik Streidl 2024-03-22 12:48:48 +01:00
parent fc4e762b05
commit fdef2abdfb
4 changed files with 25 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

14
Dockerfile
View file

@ -36,7 +36,10 @@ ARG INCLUDE_OLLAMA
## Basis ##
ENV ENV=prod \
PORT=8080 \
INCLUDE_OLLAMA_ENV=${INCLUDE_OLLAMA}
# pass build args to the build
INCLUDE_OLLAMA_DOCKER=${INCLUDE_OLLAMA} \
USE_MPS_DOCKER=${USE_MPS} \
USE_CUDA_DOCKER=${USE_CUDA}
## Basis URL Config ##
ENV OLLAMA_BASE_URL="/ollama" \
@ -65,7 +68,7 @@ ENV RAG_EMBEDDING_MODEL="all-MiniLM-L6-v2" \
# Important:
# If you want to use CUDA you need to install the nvidia-container-toolkit (https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
# you can set this to "cuda" but its recomended to use --build-arg CUDA_ENABLED=true flag when building the image
RAG_EMBEDDING_MODEL_DEVICE_TYPE="cpu" \
# RAG_EMBEDDING_MODEL_DEVICE_TYPE="cpu" \
DEVICE_COMPUTE_TYPE="int8"
# device type for whisper tts and embbeding models - "cpu" (default), "cuda" (nvidia gpu and CUDA required) or "mps" (apple silicon) - choosing this right can lead to better performance
#### Preloaded models ##########################################################
@ -75,21 +78,18 @@ WORKDIR /app/backend
COPY ./backend/requirements.txt ./requirements.txt
RUN if [ "$USE_CUDA" = "true" ]; then \
export DEVICE_TYPE="cuda" && \
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117 --no-cache-dir && \
pip3 install -r requirements.txt --no-cache-dir; \
elif [ "$USE_MPS" = "true" ]; then \
export DEVICE_TYPE="mps" && \
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu --no-cache-dir && \
pip3 install -r requirements.txt --no-cache-dir && \
python -c "import os; from faster_whisper import WhisperModel; WhisperModel(os.environ['WHISPER_MODEL'], device='cpu', compute_type='int8', download_root=os.environ['WHISPER_MODEL_DIR'])" && \
python -c "import os; from chromadb.utils import embedding_functions; sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=os.environ['RAG_EMBEDDING_MODEL'], device=os.environ['DEVICE_TYPE'])"; \
python -c "import os; from chromadb.utils import embedding_functions; sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=os.environ['RAG_EMBEDDING_MODEL'], device='mps')"; \
else \
export DEVICE_TYPE="cpu" && \
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu --no-cache-dir && \
pip3 install -r requirements.txt --no-cache-dir && \
python -c "import os; from faster_whisper import WhisperModel; WhisperModel(os.environ['WHISPER_MODEL'], device='cpu', compute_type='int8', download_root=os.environ['WHISPER_MODEL_DIR'])" && \
python -c "import os; from chromadb.utils import embedding_functions; sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=os.environ['RAG_EMBEDDING_MODEL'], device=os.environ['DEVICE_TYPE'])"; \
python -c "import os; from chromadb.utils import embedding_functions; sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=os.environ['RAG_EMBEDDING_MODEL'], device='cpu')"; \
fi