Prep for runpod docker

docker/Dockerfile

@@ -0,0 +1,21 @@
+FROM runpod/base:0.6.2-cuda12.1.0
+LABEL authors="jaret"
+
+# Install dependencies
+RUN apt-get update
+
+WORKDIR /app
+ARG CACHEBUST=1
+RUN git clone https://github.com/ostris/ai-toolkit.git && \
+    cd ai-toolkit && \
+    git submodule update --init --recursive
+
+WORKDIR /app/ai-toolkit
+
+RUN ln -s /usr/bin/python3 /usr/bin/python
+RUN python -m pip install -r requirements.txt
+
+RUN apt-get install -y tmux nvtop htop
+
+WORKDIR /
+CMD ["/start.sh"]

requirements.txt

@@ -26,4 +26,5 @@ bitsandbytes
 hf_transfer
 lpips
 pytorch_fid
-optimum-quanto
+optimum-quanto
+sentencepiece

run.py

@@ -1,5 +1,5 @@
 import os
-# os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 import sys
 from typing import Union, OrderedDict
 from dotenv import load_dotenv