From 76bf0b166b992a208c53f5cb98bd20a7e3c7f933 Mon Sep 17 00:00:00 2001
From: DekitaRPG <dekitarpg@gmail.com>
Date: Mon, 16 Oct 2023 12:26:27 +0100
Subject: [PATCH] feat: base64 image output (#8)

* base64 image output

added changes for base64 image output when no s3 bucket is set. also added basic 'custom' directory for adding models/custom nodes for being mounted, updated doc, etc. :)

* pr alterations
---
 Dockerfile        |  4 ++--
 README.md         |  3 +--
 src/rp_handler.py | 27 +++++++++++++++++++++------
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index b3292ef..1f85df1 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -32,11 +32,11 @@ RUN pip3 install --no-cache-dir torch torchvision torchaudio --index-url https:/
 # Install runpod
 RUN pip3 install runpod requests
 
-# Download the models
+# Download models to include in image. (not required if including other models below)
 RUN wget -O models/checkpoints/sd_xl_base_1.0.safetensors https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors
 RUN wget -O models/checkpoints/sdxl_vae.safetensors https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors
 
-# # Add the models
+# # Example for adding specific models into image
 # ADD models/checkpoints/sd_xl_base_1.0.safetensors models/checkpoints/
 # ADD models/checkpoints/sdxl_vae.safetensors models/checkpoints/
 
diff --git a/README.md b/README.md
index 857c12d..24df883 100644
--- a/README.md
+++ b/README.md
@@ -40,7 +40,7 @@
 ## Features
 
 * Run any [ComfyUI](https://github.com/comfyanonymous/ComfyUI) workflow to generate an image
-* Generated images are uploaded to AWS S3
+* Generated images are uploaded to AWS S3 (or returned as base64 if no bucket set)
 * Build-in checkpoint: [sd_xl_base_1.0.safetensors](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
 * Build-in VAE: [sdxl_vae.safetensors](https://huggingface.co/stabilityai/sdxl-vae/)
 * Based on [Ubuntu + NVIDIA CUDA](https://hub.docker.com/r/nvidia/cuda)
@@ -167,7 +167,6 @@ To run the Docker image on Windows, we need to have WSL2 and a Linux distro (lik
 
 - Add your user to the `docker` group, so that you can use Docker without `sudo`: `sudo usermod -aG docker $USER`
 
-
 ## Automatically deploy to Docker hub with Github Actions
 
 The repo contains two workflows that publishes the image to Docker hub using Github Actions:
diff --git a/src/rp_handler.py b/src/rp_handler.py
index b048802..c465378 100644
--- a/src/rp_handler.py
+++ b/src/rp_handler.py
@@ -6,6 +6,7 @@
 import time
 import os
 import requests
+import base64
 
 # Time to wait between API check attempts in milliseconds
 COMFY_API_AVAILABLE_INTERVAL_MS = 50
@@ -83,6 +84,15 @@ def get_history(prompt_id):
         return json.loads(response.read())
 
 
+def base64_encode(img_path):
+    """
+    Returns base64 encoded image.
+    """
+    with open(img_path, "rb") as image_file:
+        encoded_string = base64.b64encode(image_file.read())
+        return encoded_string.decode("utf-8")
+    
+
 def handler(job):
     """
     The main function that handles a job of generating an image.
@@ -160,18 +170,23 @@ def handler(job):
 
     print(f"runpod-worker-comfy - image generation is done")
 
+    # expected image output folder
+    local_image_path = f"{COMFY_OUTPUT_PATH}/{output_images}"
     # The image is in the output folder
-    if os.path.exists(f"{COMFY_OUTPUT_PATH}/{output_images}"):
+    if os.path.exists(local_image_path):
         print("runpod-worker-comfy - the image exists in the output folder")
-        image_url = rp_upload.upload_image(
-            job["id"], f"{COMFY_OUTPUT_PATH}/{output_images}"
-        )
-        return {"status": "success", "message": f"{image_url}"}
+        image_url = rp_upload.upload_image(job["id"], local_image_path)
+        return_base64 = "simulated_uploaded/" in image_url
+        return_output = f"{image_url}" if not return_base64 else base64_encode(local_image_path)
+        return {
+            "status": "success", 
+            "message": return_output, 
+        }
     else:
         print("runpod-worker-comfy - the image does not exist in the output folder")
         return {
             "status": "error",
-            "message": f"the image does not exist in the specified output folder: {COMFY_OUTPUT_PATH}/{output_images}",
+            "message": f"the image does not exist in the specified output folder: {local_image_path}",
         }