From 1d5f387f544a2d0b514f8e660cb29269f065b399 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Thu, 27 Mar 2025 17:44:46 -0600
Subject: [PATCH] Fix docker command to work better with runpod

---
 docker/Dockerfile            |  7 +++-
 docker/start.sh              | 70 ++++++++++++++++++++++++++++++++++++
 toolkit/dataloader_mixins.py |  2 +-
 version.py                   |  2 +-
 4 files changed, 78 insertions(+), 3 deletions(-)
 create mode 100644 docker/start.sh

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 99435ec..3eda6ea 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -64,4 +64,9 @@ RUN npm install && \
 # Expose port (assuming the application runs on port 3000)
 EXPOSE 8675
 
-CMD ["npm", "run", "start"]
\ No newline at end of file
+WORKDIR /
+
+COPY docker/start.sh /start.sh
+RUN chmod +x /start.sh
+
+CMD ["/start.sh"]
\ No newline at end of file
diff --git a/docker/start.sh b/docker/start.sh
new file mode 100644
index 0000000..abcbca0
--- /dev/null
+++ b/docker/start.sh
@@ -0,0 +1,70 @@
+#!/bin/bash
+set -e  # Exit the script if any statement returns a non-true return value
+
+# ref https://github.com/runpod/containers/blob/main/container-template/start.sh
+
+# ---------------------------------------------------------------------------- #
+#                          Function Definitions                                #
+# ---------------------------------------------------------------------------- #
+
+
+# Setup ssh
+setup_ssh() {
+    if [[ $PUBLIC_KEY ]]; then
+        echo "Setting up SSH..."
+        mkdir -p ~/.ssh
+        echo "$PUBLIC_KEY" >> ~/.ssh/authorized_keys
+        chmod 700 -R ~/.ssh
+
+         if [ ! -f /etc/ssh/ssh_host_rsa_key ]; then
+            ssh-keygen -t rsa -f /etc/ssh/ssh_host_rsa_key -q -N ''
+            echo "RSA key fingerprint:"
+            ssh-keygen -lf /etc/ssh/ssh_host_rsa_key.pub
+        fi
+
+        if [ ! -f /etc/ssh/ssh_host_dsa_key ]; then
+            ssh-keygen -t dsa -f /etc/ssh/ssh_host_dsa_key -q -N ''
+            echo "DSA key fingerprint:"
+            ssh-keygen -lf /etc/ssh/ssh_host_dsa_key.pub
+        fi
+
+        if [ ! -f /etc/ssh/ssh_host_ecdsa_key ]; then
+            ssh-keygen -t ecdsa -f /etc/ssh/ssh_host_ecdsa_key -q -N ''
+            echo "ECDSA key fingerprint:"
+            ssh-keygen -lf /etc/ssh/ssh_host_ecdsa_key.pub
+        fi
+
+        if [ ! -f /etc/ssh/ssh_host_ed25519_key ]; then
+            ssh-keygen -t ed25519 -f /etc/ssh/ssh_host_ed25519_key -q -N ''
+            echo "ED25519 key fingerprint:"
+            ssh-keygen -lf /etc/ssh/ssh_host_ed25519_key.pub
+        fi
+
+        service ssh start
+
+        echo "SSH host keys:"
+        for key in /etc/ssh/*.pub; do
+            echo "Key: $key"
+            ssh-keygen -lf $key
+        done
+    fi
+}
+
+# Export env vars
+export_env_vars() {
+    echo "Exporting environment variables..."
+    printenv | grep -E '^RUNPOD_|^PATH=|^_=' | awk -F = '{ print "export " $1 "=\"" $2 "\"" }' >> /etc/rp_environment
+    echo 'source /etc/rp_environment' >> ~/.bashrc
+}
+
+# ---------------------------------------------------------------------------- #
+#                               Main Program                                   #
+# ---------------------------------------------------------------------------- #
+
+
+echo "Pod Started"
+
+setup_ssh
+export_env_vars
+echo "Starting AI Toolkit UI..."
+cd /app/ai-toolkit/ui && npm run start 
\ No newline at end of file
diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py
index 115a6ac..8649129 100644
--- a/toolkit/dataloader_mixins.py
+++ b/toolkit/dataloader_mixins.py
@@ -771,7 +771,7 @@ class InpaintControlFileItemDTOMixin:
             img = Image.open(self.inpaint_path)
             # make sure has aplha
             if img.mode != 'RGBA':
-                raise ValueError(f"Image must have alpha channel for inpaint: {self.inpaint_path}")
+                return
             img = exif_transpose(img)
         
             w, h = img.size
diff --git a/version.py b/version.py
index 826c403..14c32f0 100644
--- a/version.py
+++ b/version.py
@@ -1 +1 @@
-VERSION = "0.2.3"
\ No newline at end of file
+VERSION = "0.2.4"
\ No newline at end of file