From aa21115e26ea3cdc20e3674ba3b3e98a9da8d113 Mon Sep 17 00:00:00 2001
From: Marc-Andre Ferland <madrang@gmail.com>
Date: Tue, 1 Nov 2022 04:52:42 -0400
Subject: [PATCH] Always return a byte buffer. Sending the picture as URL text
 fails in some browsers.

---
 ui/sd_internal/runtime.py      | 6 +++++-
 ui/sd_internal/task_manager.py | 3 ++-
 ui/server.py                   | 2 --
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/ui/sd_internal/runtime.py b/ui/sd_internal/runtime.py
index ec2fe993..847ddb79 100644
--- a/ui/sd_internal/runtime.py
+++ b/ui/sd_internal/runtime.py
@@ -870,9 +870,13 @@ def img_to_base64_str(img, output_format="PNG"):
     img_str = "data:image/png;base64," + base64.b64encode(img_byte).decode()
     return img_str
 
-def base64_str_to_img(img_str):
+def base64_str_to_buffer(img_str):
     img_str = img_str[len("data:image/png;base64,"):]
     data = base64.b64decode(img_str)
     buffered = BytesIO(data)
+    return buffered
+
+def base64_str_to_img(img_str):
+    buffered = base64_str_to_buffer(img_str)
     img = Image.open(buffered)
     return img
diff --git a/ui/sd_internal/task_manager.py b/ui/sd_internal/task_manager.py
index 49e6addf..822379fc 100644
--- a/ui/sd_internal/task_manager.py
+++ b/ui/sd_internal/task_manager.py
@@ -334,7 +334,8 @@ def thread_render(device):
                             img_id = out_obj['path'][out_obj['path'].rindex('/') + 1:]
                             task.temp_images[int(img_id)] = runtime.thread_data.temp_images[out_obj['path'][11:]]
                         elif 'data' in out_obj:
-                            task.temp_images[result['output'].index(out_obj)] = out_obj['data']
+                            buf = runtime.base64_str_to_buffer(out_obj['data'])
+                            task.temp_images[result['output'].index(out_obj)] = buf
                 # Before looping back to the generator, mark cache as still alive.
                 task_cache.keep(task.request.session_id, TASK_TTL)
         except Exception as e:
diff --git a/ui/server.py b/ui/server.py
index 57d374b0..5ed5fb5d 100644
--- a/ui/server.py
+++ b/ui/server.py
@@ -336,8 +336,6 @@ def get_image(session_id, img_id):
     if not task.temp_images[img_id]: raise HTTPException(status_code=425, detail='Too Early, task data is not available yet.') # HTTP425 Too Early
     try:
         img_data = task.temp_images[img_id]
-        if isinstance(img_data, str):
-            return img_data
         img_data.seek(0)
         return StreamingResponse(img_data, media_type='image/jpeg')
     except KeyError as e: