Gallery improvements (#4789)

extensions/gallery/script.js

@@ -5,6 +5,13 @@ let extensions_block = document.getElementById('extensions');
 let extensions_block_size = extensions_block.childNodes.length;
 let gallery_only = (extensions_block_size == 5);
 
+function gotoFirstPage() {
+    const firstPageButton = gallery_element.querySelector('.paginate > button');
+    if (firstPageButton) {
+        firstPageButton.click();
+    }
+}
+
 document.querySelector('.header_bar').addEventListener('click', function(event) {
     if (event.target.tagName === 'BUTTON') {
         const buttonText = event.target.textContent.trim();

extensions/gallery/script.py

@@ -3,11 +3,17 @@ from pathlib import Path
 import gradio as gr
 
 from modules.html_generator import get_image_cache
-from modules.shared import gradio
+from modules.shared import gradio, settings
+
+cards = []
 
 
 def generate_css():
     css = """
+      .highlighted-border {
+        border-color: rgb(249, 115, 22) !important;
+      }
+
       .character-gallery > .gallery {
         margin: 1rem 0;
         display: grid !important;
@@ -58,6 +64,7 @@ def generate_css():
 
 
 def generate_html():
+    global cards
     cards = []
     # Iterate through files in image folder
     for file in sorted(Path("characters").glob("*")):
@@ -78,6 +85,14 @@ def generate_html():
     return cards
 
 
+def filter_cards(filter_str=''):
+    if filter_str == '':
+        return cards
+
+    filter_upper = filter_str.upper()
+    return [k for k in cards if filter_upper in k[1].upper()]
+
+
 def select_character(evt: gr.SelectData):
     return (evt.value[1])
 
@@ -88,16 +103,25 @@ def custom_js():
 
 
 def ui():
-    with gr.Accordion("Character gallery", open=False, elem_id='gallery-extension'):
-        update = gr.Button("Refresh")
+    with gr.Accordion("Character gallery", open=settings["gallery-open"], elem_id='gallery-extension'):
         gr.HTML(value="<style>" + generate_css() + "</style>")
+        with gr.Row():
+            filter_box = gr.Textbox(label='', placeholder='Filter', lines=1, max_lines=1, container=False, elem_id='gallery-filter-box')
+            gr.ClearButton(filter_box, value='🗑️', elem_classes='refresh-button')
+            update = gr.Button("Refresh", elem_classes='refresh-button')
         gallery = gr.Dataset(
             components=[gr.HTML(visible=False)],
             label="",
             samples=generate_html(),
             elem_classes=["character-gallery"],
-            samples_per_page=50
+            samples_per_page=settings["gallery-items_per_page"]
         )
 
-    update.click(generate_html, [], gallery)
+    filter_box.change(lambda: None, None, None, _js=f'() => {{{custom_js()}; gotoFirstPage()}}').success(
+        filter_cards, filter_box, gallery).then(
+        lambda x: gr.update(elem_classes='highlighted-border' if x != '' else ''), filter_box, filter_box, show_progress=False)
+
+    update.click(generate_html, [], None).success(
+        filter_cards, filter_box, gallery)
+
     gallery.select(select_character, None, gradio['character_menu'])

modules/chat.py

@@ -281,7 +281,7 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess
 
 def impersonate_wrapper(text, state):
 
-    static_output = chat_html_wrapper(state['history'], state['name1'], state['name2'], state['mode'], state['chat_style'])
+    static_output = chat_html_wrapper(state['history'], state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu'])
 
     if shared.model_name == 'None' or shared.model is None:
         logger.error("No model is loaded! Select one in the Model tab.")
@@ -340,7 +340,7 @@ def generate_chat_reply_wrapper(text, state, regenerate=False, _continue=False):
         send_dummy_reply(state['start_with'], state)
 
     for i, history in enumerate(generate_chat_reply(text, state, regenerate, _continue, loading_message=True)):
-        yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style']), history
+        yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu']), history
 
 
 def remove_last_message(history):
@@ -390,8 +390,8 @@ def send_dummy_reply(text, state):
     return history
 
 
-def redraw_html(history, name1, name2, mode, style, reset_cache=False):
-    return chat_html_wrapper(history, name1, name2, mode, style, reset_cache=reset_cache)
+def redraw_html(history, name1, name2, mode, style, character, reset_cache=False):
+    return chat_html_wrapper(history, name1, name2, mode, style, character, reset_cache=reset_cache)
 
 
 def start_new_chat(state):

modules/html_generator.py

@@ -221,11 +221,11 @@ def generate_instruct_html(history):
     return output
 
 
-def generate_cai_chat_html(history, name1, name2, style, reset_cache=False):
+def generate_cai_chat_html(history, name1, name2, style, character, reset_cache=False):
     output = f'<style>{chat_styles[style]}</style><div class="chat" id="chat"><div class="messages">'
 
-    # We use ?name2 and ?time.time() to force the browser to reset caches
-    img_bot = f'<img src="file/cache/pfp_character_thumb.png?{name2}" class="pfp_character">' if Path("cache/pfp_character_thumb.png").exists() else ''
+    # We use ?character and ?time.time() to force the browser to reset caches
+    img_bot = f'<img src="file/cache/pfp_character_thumb.png?{character}" class="pfp_character">' if Path("cache/pfp_character_thumb.png").exists() else ''
     img_me = f'<img src="file/cache/pfp_me.png?{time.time() if reset_cache else ""}">' if Path("cache/pfp_me.png").exists() else ''
 
     for i, _row in enumerate(history):
@@ -299,10 +299,10 @@ def generate_chat_html(history, name1, name2, reset_cache=False):
     return output
 
 
-def chat_html_wrapper(history, name1, name2, mode, style, reset_cache=False):
+def chat_html_wrapper(history, name1, name2, mode, style, character, reset_cache=False):
     if mode == 'instruct':
         return generate_instruct_html(history['visible'])
     elif style == 'wpp':
         return generate_chat_html(history['visible'], name1, name2)
     else:
-        return generate_cai_chat_html(history['visible'], name1, name2, style, reset_cache)
+        return generate_cai_chat_html(history['visible'], name1, name2, style, character, reset_cache)

modules/shared.py

@@ -58,6 +58,8 @@ settings = {
     'custom_system_message': '',
     'chat-instruct_command': 'Continue the chat dialogue below. Write a single reply for the character "<|character|>".\n\n<|prompt|>',
     'autoload_model': False,
+    'gallery-items_per_page': 50,
+    'gallery-open': False,
     'default_extensions': ['gallery'],
 }
 

modules/ui_chat.py

@@ -11,7 +11,7 @@ from modules.text_generation import stop_everything_event
 from modules.utils import gradio
 
 inputs = ('Chat input', 'interface_state')
-reload_arr = ('history', 'name1', 'name2', 'mode', 'chat_style')
+reload_arr = ('history', 'name1', 'name2', 'mode', 'chat_style', 'character_menu')
 clear_arr = ('delete_chat-confirm', 'delete_chat', 'delete_chat-cancel')
 
 
@@ -25,7 +25,7 @@ def create_ui():
     with gr.Tab('Chat', elem_id='chat-tab', elem_classes=("old-ui" if shared.args.chat_buttons else None)):
         with gr.Row():
             with gr.Column(elem_id='chat-col'):
-                shared.gradio['display'] = gr.HTML(value=chat_html_wrapper({'internal': [], 'visible': []}, '', '', 'chat', 'cai-chat'))
+                shared.gradio['display'] = gr.HTML(value=chat_html_wrapper({'internal': [], 'visible': []}, '', '', 'chat', 'cai-chat', ''))
 
                 with gr.Row(elem_id="chat-input-row"):
                     with gr.Column(scale=1, elem_id='gr-hover-container'):

settings-template.yaml

@@ -30,5 +30,7 @@ chat-instruct_command: |-
 
   <|prompt|>
 autoload_model: false
+gallery-items_per_page: 50
+gallery-open: false
 default_extensions:
 - gallery