From 6dac5341dbf765d346da891ddef3982591a5a5fa Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Wed, 24 Aug 2022 11:18:25 +0530
Subject: [PATCH] Can feed a custom image

---
 server.sh | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/server.sh b/server.sh
index b91d095b..40f23546 100644
--- a/server.sh
+++ b/server.sh
@@ -1,7 +1,12 @@
 #!/bin/bash
 
-IMAGE_NAME="r8.im/stability-ai/stable-diffusion@sha256:06eb78b36068500c616a7f33c15e6fa40404f8e14b5bfad57ebe0c7fe0f6bdf1"
+IMAGE_NAME=$1
 
+if [ -z "$IMAGE_NAME" ]; then
+  IMAGE_NAME="r8.im/stability-ai/stable-diffusion@sha256:06eb78b36068500c616a7f33c15e6fa40404f8e14b5bfad57ebe0c7fe0f6bdf1"
+fi
+
+echo "Using docker image $IMAGE_NAME"
 docker run --name sd -d -p 5000:5000 --gpus all $IMAGE_NAME
 
 uvicorn main:app --reload
\ No newline at end of file