From 43b6ab8673b570420c784cb419f1bf5ad824286b Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Mon, 27 Feb 2023 13:41:00 -0300
Subject: [PATCH] Store thumbnails as files instead of base64 strings

This improves the UI responsiveness for large histories.
---
 extensions/gallery/script.py |  4 ++--
 modules/html_generator.py    | 21 +++++++++++----------
 server.py                    |  4 ++--
 3 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/extensions/gallery/script.py b/extensions/gallery/script.py
index 98d66343..bec26a48 100644
--- a/extensions/gallery/script.py
+++ b/extensions/gallery/script.py
@@ -2,7 +2,7 @@ from pathlib import Path
 
 import gradio as gr
 
-from modules.html_generator import image_to_base64
+from modules.html_generator import get_image_cache
 
 
 def generate_html():
@@ -64,7 +64,7 @@ def generate_html():
                 path = Path(i)
                 if path.exists():
                     try:
-                        image_html = f'<img src="data:image/png;base64,{image_to_base64(path)}">'
+                        image_html = f'<img src="file/{get_image_cache(path)}">'
                         break
                     except:
                         continue
diff --git a/modules/html_generator.py b/modules/html_generator.py
index dc2aa387..162040ba 100644
--- a/modules/html_generator.py
+++ b/modules/html_generator.py
@@ -4,15 +4,13 @@ This is a library for formatting GPT-4chan and chat outputs as nice HTML.
 
 '''
 
-import base64
 import os
 import re
-from io import BytesIO
 from pathlib import Path
 
 from PIL import Image
 
-# This is to store chat profile pictures as base64-encoded thumbnails
+# This is to store the paths to the thumbnails of the profile pictures
 image_cache = {}
 
 def generate_basic_html(s):
@@ -195,15 +193,18 @@ def generate_4chan_html(f):
 
     return output
 
-def image_to_base64(path):
-    mtime = os.stat(path).st_mtime
+def get_image_cache(path):
+    cache_folder = Path("cache")
+    if not cache_folder.exists():
+        cache_folder.mkdir()
 
+    mtime = os.stat(path).st_mtime
     if (path in image_cache and mtime != image_cache[path][0]) or (path not in image_cache):
         img = Image.open(path)
         img.thumbnail((200, 200))
-        img_buffer = BytesIO()
-        img.convert('RGB').save(img_buffer, format='PNG')
-        image_cache[path] = [mtime, base64.b64encode(img_buffer.getvalue()).decode("utf-8")]
+        output_file = Path(f'cache/{path.name}_cache.png')
+        img.convert('RGB').save(output_file, format='PNG')
+        image_cache[path] = [mtime, output_file.as_posix()]
 
     return image_cache[path][1]
 
@@ -301,14 +302,14 @@ def generate_chat_html(history, name1, name2, character):
 
         path = Path(i)
         if path.exists():
-            img = f'<img src="data:image/png;base64,{image_to_base64(path)}">'
+            img = f'<img src="file/{get_image_cache(path)}">'
             break
 
     img_me = ''
     for i in ["img_me.png", "img_me.jpg", "img_me.jpeg"]:
         path = Path(i)
         if path.exists():
-            img_me = f'<img src="data:image/png;base64,{image_to_base64(path)}">'
+            img_me = f'<img src="file/{get_image_cache(path)}">'
             break
 
     for i,_row in enumerate(history[::-1]):
diff --git a/server.py b/server.py
index 37d290a3..5d49cb47 100644
--- a/server.py
+++ b/server.py
@@ -267,8 +267,8 @@ if shared.args.chat or shared.args.cai_chat:
         # Clear history with confirmation
         clear_arr = [shared.gradio[k] for k in ['Clear history-confirm', 'Clear history', 'Clear history-cancel']]
         shared.gradio['Clear history'].click(lambda :[gr.update(visible=True), gr.update(visible=False), gr.update(visible=True)], None, clear_arr)
-        shared.gradio['Clear history-confirm'].click(chat.clear_chat_log, [shared.gradio['name1'], shared.gradio['name2']], shared.gradio['display'])
         shared.gradio['Clear history-confirm'].click(lambda :[gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)], None, clear_arr)
+        shared.gradio['Clear history-confirm'].click(chat.clear_chat_log, [shared.gradio['name1'], shared.gradio['name2']], shared.gradio['display'])
         shared.gradio['Clear history-cancel'].click(lambda :[gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)], None, clear_arr)
 
         shared.gradio['Remove last'].click(chat.remove_last_message, [shared.gradio['name1'], shared.gradio['name2']], [shared.gradio['display'], shared.gradio['textbox']], show_progress=False)
@@ -279,7 +279,7 @@ if shared.args.chat or shared.args.cai_chat:
         for i in ['Generate', 'Regenerate', 'Replace last reply']:
             shared.gradio[i].click(lambda x: '', shared.gradio['textbox'], shared.gradio['textbox'], show_progress=False)
             shared.gradio[i].click(lambda : chat.save_history(timestamp=False), [], [], show_progress=False)
-        shared.gradio['Clear history'].click(lambda : chat.save_history(timestamp=False), [], [], show_progress=False)
+        shared.gradio['Clear history-confirm'].click(lambda : chat.save_history(timestamp=False), [], [], show_progress=False)
         shared.gradio['textbox'].submit(lambda x: '', shared.gradio['textbox'], shared.gradio['textbox'], show_progress=False)
         shared.gradio['textbox'].submit(lambda : chat.save_history(timestamp=False), [], [], show_progress=False)