Install correct version of Torch

Dockerfile

@@ -57,17 +57,23 @@ RUN apt update && \
 # Set Python
 RUN ln -s /usr/bin/python3.10 /usr/bin/python
 
-# Stage 2: Install FaceFusion and python modules
+# Stage 2: Install LLaVA and python modules
 FROM base as setup
 
 # Create and use the Python venv
 RUN python3 -m venv /venv
 
-# Clone the git repo of FaceFusion and set version
+# Install Torch
+RUN source /venv/bin/activate && \
+    pip3 install --no-cache-dir torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 && \
+    pip3 install --no-cache-dir xformers==0.0.22 && \
+    deactivate
+
+# Clone the git repo of LLaVA and set version
 WORKDIR /
-RUN git clone https://github.com/haotian-liu/LLaVA.git && \
-    cd /LLaVA && \
-    git checkout ${LLAVA_VERSION}
+RUN git clone https://github.com/ashleykleynhans/LLaVA.git
+#    cd /LLaVA && \
+#    git checkout ${LLAVA_VERSION}
 
 # Install the dependencies for LLaVA
 WORKDIR /LLaVA
@@ -77,6 +83,7 @@ RUN source /venv/bin/activate && \
     pip3 install -e . && \
     pip3 install ninja && \
     pip3 install flash-attn --no-build-isolation && \
+    pip3 install transformers==4.34.1 && \
     deactivate
 
 # Install Jupyter

README.md

@@ -1,10 +1,5 @@
 # Docker image for LLaVA: Large Language and Vision Assistant
 
-> [!IMPORTANT]  
-> If you are using the 13b model, CUDA will result in OOM errors
-> with a GPU that has less than 48GB of VRAM, so A6000 or higher is
-> recommended.
-
 ## Installs
 
 * Ubuntu 22.04 LTS
@@ -12,7 +7,8 @@
 * Python 3.10.12
 * [LLaVA](
   https://github.com/haotian-liu/llava) v1.1.1
-* Torch 2.1.0
+* Torch 2.0.1
+* BakLLaVA-1 model
 
 ## Available on RunPod
 
@@ -37,11 +33,28 @@ docker run -d \
   -p 3000:3001 \
   -p 8888:8888 \
   -e JUPYTER_PASSWORD=Jup1t3R! \
-  ashleykza/llava:1.1.2
+  ashleykza/llava:latest
 ```
 
 You can obviously substitute the image name and tag with your own.
 
+#### Models
+
+> [!IMPORTANT]
+> If you select the 13b model, CUDA will result in OOM errors
+> with a GPU that has less than 48GB of VRAM, so A6000 or higher is
+> recommended.
+
+You can add an environment called `MODEL` to your Docker container to
+specify the model that should be downloaded.  If the `MODEL` environment
+variable is not set, the model will default to `SkunkworksAI/BakLLaVA-1`.
+
+| Model                                                              | Environment Variable Value  | Default |
+|--------------------------------------------------------------------|-----------------------------|---------|
+| [llava-v1.5-13b](https://huggingface.co/liuhaotian/llava-v1.5-13b) | liuhaotian/llava-v1.5-13b   | no      |
+| [llava-v1.5-7b](https://huggingface.co/liuhaotian/llava-v1.5-7b)   | liuhaotian/llava-v1.5-7b    | no      |
+| [BakLLaVA-1](https://huggingface.co/SkunkworksAI/BakLLaVA-1)       | SkunkworksAI/BakLLaVA-1     | yes     |
+
 ## Acknowledgements
 
 1. Matthew Berman for giving me a demo on LLaVA, as well as his amazing

scripts/pre_start.sh

@@ -31,7 +31,7 @@ else
     then
       export LLAVA_MODEL=${MODEL}
     else
-      export LLAVA_MODEL="liuhaotian/llava-v1.5-7b"
+      export LLAVA_MODEL="SkunkworksAI/BakLLaVA-1"
     fi
 
     mkdir -p /workspace/logs