Store thumbnails as files instead of base64 strings

This improves the UI responsiveness for large histories.
This commit is contained in:
oobabooga 2023-02-27 13:41:00 -03:00
parent a08802bf70
commit 43b6ab8673
3 changed files with 15 additions and 14 deletions

View File

@ -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

View File

@ -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]):

View File

@ -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)