From 88acc28d7f9550926994132d9eda3717e57ca4e3 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Sat, 3 Aug 2024 12:41:06 -0600
Subject: [PATCH] Prep for runpod docker

---
 docker/Dockerfile | 21 +++++++++++++++++++++
 requirements.txt  |  3 ++-
 run.py            |  2 +-
 3 files changed, 24 insertions(+), 2 deletions(-)
 create mode 100644 docker/Dockerfile

diff --git a/docker/Dockerfile b/docker/Dockerfile
new file mode 100644
index 00000000..0c548507
--- /dev/null
+++ b/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"]
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index d51fb25c..23781266 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -26,4 +26,5 @@ bitsandbytes
 hf_transfer
 lpips
 pytorch_fid
-optimum-quanto
\ No newline at end of file
+optimum-quanto
+sentencepiece
\ No newline at end of file
diff --git a/run.py b/run.py
index 068abe9a..6f133081 100644
--- a/run.py
+++ b/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