From 34fa6e38e73d3bb62ce9179f321fbe8d47b5d558 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Sat, 17 Sep 2022 08:39:20 -0400
Subject: [PATCH] fix long hex-encoded error message from legacy server

- closes issue #618
---
 ldm/dream/args.py   | 20 +++++++++++++++++++-
 ldm/dream/server.py |  1 -
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/ldm/dream/args.py b/ldm/dream/args.py
index e9b70a7199..626d66c91c 100644
--- a/ldm/dream/args.py
+++ b/ldm/dream/args.py
@@ -61,6 +61,7 @@ import json
 import hashlib
 import os
 import copy
+import base64
 from ldm.dream.conditioning import split_weighted_subprompts
 
 SAMPLER_CHOICES = [
@@ -593,7 +594,7 @@ def format_metadata(opt,
     if opt.init_img:
         rfc_dict['type']           = 'img2img'
         rfc_dict['strength_steps'] = rfc_dict.pop('strength')
-        rfc_dict['orig_hash']      = sha256(image_dict['init_img'])
+        rfc_dict['orig_hash']      = calculate_init_img_hash(opt.init_img)
         rfc_dict['sampler']        = 'ddim'  # FIX ME WHEN IMG2IMG SUPPORTS ALL SAMPLERS
     else:
         rfc_dict['type']  = 'txt2img'
@@ -612,6 +613,23 @@ def format_metadata(opt,
         'images'      : images,
     }
 
+# image can either be a file path on disk or a base64-encoded
+# representation of the file's contents
+def calculate_init_img_hash(image_string):
+    prefix = 'data:image/png;base64,'
+    hash   = None
+    if image_string.startswith(prefix):
+        imagebase64 = image_string[len(prefix):]
+        imagedata   = base64.b64decode(imagebase64)
+        with open('outputs/test.png','wb') as file:
+            file.write(imagedata)
+        sha = hashlib.sha256()
+        sha.update(imagedata)
+        hash = sha.hexdigest()
+    else:
+        hash = sha256(image_string)
+    return hash
+
 # Bah. This should be moved somewhere else...
 def sha256(path):
     sha = hashlib.sha256()
diff --git a/ldm/dream/server.py b/ldm/dream/server.py
index 9147a3180a..5dd8f38c2f 100644
--- a/ldm/dream/server.py
+++ b/ldm/dream/server.py
@@ -145,7 +145,6 @@ class DreamServer(BaseHTTPRequestHandler):
         opt = build_opt(post_data, self.model.seed, gfpgan_model_exists)
 
         self.canceled.clear()
-        print(f">> Request to generate with prompt: {opt.prompt}")
         # In order to handle upscaled images, the PngWriter needs to maintain state
         # across images generated by each call to prompt2img(), so we define it in
         # the outer scope of image_done()