From 48def808126209bb3e67b8fdb53a34820a8a0665 Mon Sep 17 00:00:00 2001
From: Hendrik Langer <hendrik+dev@xd0.de>
Date: Wed, 5 Apr 2023 18:42:01 +0200
Subject: [PATCH] add worker container

---
 runpod/serverless-automatic/Dockerfile | 37 +++++++++++++++++++++
 runpod/serverless-automatic/handler.py | 45 ++++++++++++++++++++++++++
 runpod/serverless-automatic/start.sh   | 11 +++++++
 3 files changed, 93 insertions(+)
 create mode 100644 runpod/serverless-automatic/Dockerfile
 create mode 100644 runpod/serverless-automatic/handler.py
 create mode 100644 runpod/serverless-automatic/start.sh

diff --git a/runpod/serverless-automatic/Dockerfile b/runpod/serverless-automatic/Dockerfile
new file mode 100644
index 0000000..e4f3cac
--- /dev/null
+++ b/runpod/serverless-automatic/Dockerfile
@@ -0,0 +1,37 @@
+# https://blog.runpod.io/custom-models-with-serverless-and-automatic-stable-diffusion/
+FROM runpod/stable-diffusion:web-automatic-base-4.0.0
+
+SHELL ["/bin/bash", "-c"]
+
+ENV PATH="${PATH}:/workspace/stable-diffusion-webui/venv/bin"
+
+WORKDIR /workspace
+
+# PFG: https://civitai.com/api/download/models/1316
+# Hassanblend: https://civitai.com/api/download/models/4635
+# Deliberate v2: https://civitai.com/api/download/models/15236
+
+RUN wget -O /workspace/stable-diffusion-webui/models/Stable-diffusion/model.safetensors https://civitai.com/api/download/models/4635
+## Extra downloads (for Hassanblend)
+RUN wget -O /workspace/stable-diffusion-webui/models/VAE/vae-ft-mse-840000-ema-pruned.safetensors https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors
+
+RUN mkdir /workspace/stable-diffusion-webui/models/Codeformer/ && \
+    wget -O /workspace/stable-diffusion-webui/models/Codeformer/codeformer-v0.1.0.pth https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth && \
+    wget -O /workspace/stable-diffusion-webui/repositories/CodeFormer/weights/facelib/detection_Resnet50_Final.pth https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth && \
+    wget -O /workspace/stable-diffusion-webui/repositories/CodeFormer/weights/facelib/parsing_parsenet.pth https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth && \
+    mkdir /workspace/stable-diffusion-webui/models/GFPGAN/ && \
+    wget -O /workspace/stable-diffusion-webui/models/GFPGAN/GFPGANv1.4.pth https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth && \
+    mkdir /workspace/stable-diffusion-webui/models/RealESRGAN/ && \
+    wget -O /workspace/stable-diffusion-webui/models/RealESRGAN/RealESRGAN_x4plus.pth https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth && \
+    mkdir /workspace/stable-diffusion-webui/models/ESRGAN/ && \
+    wget -O /workspace/stable-diffusion-webui/models/ESRGAN/ESRGAN.pth https://github.com/cszn/KAIR/releases/download/v1.0/ESRGAN.pth && \
+#    mkdir /workspace/stable-diffusion-webui/models/Stable-diffusion/tokenizer/ && \
+#    wget -O /workspace/stable-diffusion-webui/models/Stable-diffusion/tokenizer/vocab.json https://huggingface.co/runwayml/stable-diffusion-v1-5/raw/main/tokenizer/vocab.json && \
+    git clone --depth 1 https://github.com/Coyote-A/ultimate-upscale-for-automatic1111 /workspace/stable-diffusion-webui/extensions/ultimate-upscale-for-automatic1111
+RUN pip install runpod
+
+ADD handler.py /workspace/
+ADD start.sh /workspace/
+RUN chmod +x /workspace/start.sh
+
+CMD [ "/workspace/start.sh" ]
diff --git a/runpod/serverless-automatic/handler.py b/runpod/serverless-automatic/handler.py
new file mode 100644
index 0000000..28d5723
--- /dev/null
+++ b/runpod/serverless-automatic/handler.py
@@ -0,0 +1,45 @@
+import runpod
+import subprocess
+import requests
+import time
+
+def check_api_availability(host):
+    while True:
+        try:
+            response = requests.get(host)
+            return
+        except requests.exceptions.RequestException as e:
+            print(f"API is not available, retrying in 200ms... ({e})")
+        except Exception as e:
+            print('something went wrong')
+        time.sleep(200/1000)
+
+check_api_availability("http://127.0.0.1:3000/sdapi/v1/txt2img")
+
+print('run handler')
+
+def handler(event):
+    '''
+    This is the handler function that will be called by the serverless.
+    '''
+    print('got event')
+    print(event)
+
+    cmd = 'txt2img'
+    if 'cmd' in event:
+        cmd = event["cmd"]
+    if 'api_endpoint' in event["input"]:
+        cmd = event["input"]["api_endpoint"]
+        del event["input"]["api_endpoint"]
+    response = requests.post(url=f'http://127.0.0.1:3000/sdapi/v1/{cmd}', json=event["input"])
+
+    json = response.json()
+    # do the things
+
+    print(json)
+
+    # return the output that you want to be returned like pre-signed URLs to output artifacts
+    return json
+
+
+runpod.serverless.start({"handler": handler})
diff --git a/runpod/serverless-automatic/start.sh b/runpod/serverless-automatic/start.sh
new file mode 100644
index 0000000..1e8951f
--- /dev/null
+++ b/runpod/serverless-automatic/start.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+echo "Container Started"
+export PYTHONUNBUFFERED=1
+source /workspace/stable-diffusion-webui/venv/bin/activate
+cd /workspace/stable-diffusion-webui
+echo "starting api"
+python webui.py --port 3000 --nowebui --api --xformers --ckpt /workspace/stable-diffusion-webui/models/Stable-diffusion/model.safetensors &
+
+echo "starting worker"
+cd /workspace
+python -u handler.py