feat: Add a dropdown to save/load user personas (#7367)

css/main.css

@@ -1645,7 +1645,7 @@ button:focus {
 }
 
 #user-description textarea {
-    height: calc(100vh - 231px) !important;
+    height: calc(100vh - 334px) !important;
     min-height: 90px !important;
 }
 

modules/chat.py

@@ -32,7 +32,12 @@ from modules.text_generation import (
     get_encoded_length,
     get_max_prompt_length
 )
-from modules.utils import delete_file, get_available_characters, save_file
+from modules.utils import (
+    delete_file,
+    get_available_characters,
+    get_available_users,
+    save_file
+)
 from modules.web_search import add_web_search_attachments
 
 
@@ -1647,6 +1652,150 @@ def delete_character(name, instruct=False):
         delete_file(Path(f'user_data/characters/{name}.{extension}'))
 
 
+def generate_user_pfp_cache(user):
+    """Generate cached profile picture for user"""
+    cache_folder = Path(shared.args.disk_cache_dir)
+    if not cache_folder.exists():
+        cache_folder.mkdir()
+
+    for path in [Path(f"user_data/users/{user}.{extension}") for extension in ['png', 'jpg', 'jpeg']]:
+        if path.exists():
+            original_img = Image.open(path)
+            # Define file paths
+            pfp_path = Path(f'{cache_folder}/pfp_me.png')
+
+            # Save thumbnail
+            thumb = make_thumbnail(original_img)
+            thumb.save(pfp_path, format='PNG')
+            logger.info(f'User profile picture cached to "{pfp_path}"')
+
+            return str(pfp_path)
+
+    return None
+
+
+def load_user(user_name, name1, user_bio):
+    """Load user profile from YAML file"""
+    picture = None
+
+    filepath = None
+    for extension in ["yml", "yaml", "json"]:
+        filepath = Path(f'user_data/users/{user_name}.{extension}')
+        if filepath.exists():
+            break
+
+    if filepath is None or not filepath.exists():
+        logger.error(f"Could not find the user \"{user_name}\" inside user_data/users. No user has been loaded.")
+        raise ValueError
+
+    with open(filepath, 'r', encoding='utf-8') as f:
+        file_contents = f.read()
+
+    extension = filepath.suffix[1:]  # Remove the leading dot
+    data = json.loads(file_contents) if extension == "json" else yaml.safe_load(file_contents)
+
+    # Clear existing user picture cache
+    cache_folder = Path(shared.args.disk_cache_dir)
+    pfp_path = Path(f"{cache_folder}/pfp_me.png")
+    if pfp_path.exists():
+        pfp_path.unlink()
+
+    # Generate new picture cache
+    picture = generate_user_pfp_cache(user_name)
+
+    # Get user name
+    if 'name' in data and data['name'] != '':
+        name1 = data['name']
+
+    # Get user bio
+    if 'user_bio' in data:
+        user_bio = data['user_bio']
+
+    return name1, user_bio, picture
+
+
+def generate_user_yaml(name, user_bio):
+    """Generate YAML content for user profile"""
+    data = {
+        'name': name,
+        'user_bio': user_bio,
+    }
+
+    return yaml.dump(data, sort_keys=False, width=float("inf"))
+
+
+def save_user(name, user_bio, picture, filename):
+    """Save user profile to YAML file"""
+    if filename == "":
+        logger.error("The filename is empty, so the user will not be saved.")
+        return
+
+    # Ensure the users directory exists
+    users_dir = Path('user_data/users')
+    users_dir.mkdir(parents=True, exist_ok=True)
+
+    data = generate_user_yaml(name, user_bio)
+    filepath = Path(f'user_data/users/{filename}.yaml')
+    save_file(filepath, data)
+
+    path_to_img = Path(f'user_data/users/{filename}.png')
+    if picture is not None:
+        # Copy the image file from its source path to the users folder
+        shutil.copy(picture, path_to_img)
+        logger.info(f'Saved user profile picture to {path_to_img}.')
+
+
+def delete_user(name):
+    """Delete user profile files"""
+    # Check for user data files
+    for extension in ["yml", "yaml", "json"]:
+        delete_file(Path(f'user_data/users/{name}.{extension}'))
+
+    # Check for user image files
+    for extension in ["png", "jpg", "jpeg"]:
+        delete_file(Path(f'user_data/users/{name}.{extension}'))
+
+
+def update_user_menu_after_deletion(idx):
+    """Update user menu after a user is deleted"""
+    users = get_available_users()
+    if len(users) == 0:
+        # Create a default user if none exist
+        save_user('You', '', None, 'Default')
+        users = get_available_users()
+
+    idx = min(int(idx), len(users) - 1)
+    idx = max(0, idx)
+    return gr.update(choices=users, value=users[idx])
+
+
+def handle_user_menu_change(state):
+    """Handle user menu selection change"""
+    try:
+        name1, user_bio, picture = load_user(state['user_menu'], state['name1'], state['user_bio'])
+
+        return [
+            name1,
+            user_bio,
+            picture
+        ]
+    except Exception as e:
+        logger.error(f"Failed to load user '{state['user_menu']}': {e}")
+        return [
+            state['name1'],
+            state['user_bio'],
+            None
+        ]
+
+
+def handle_save_user_click(name1):
+    """Handle save user button click"""
+    return [
+        name1,
+        gr.update(visible=True)
+    ]
+
+
 def jinja_template_from_old_format(params, verbose=False):
     MASTER_TEMPLATE = """
 {%- set ns = namespace(found=false) -%}

modules/shared.py

@@ -298,6 +298,7 @@ settings = {
 
     # Character settings
     'character': 'Assistant',
+    'user': 'Default',
     'name1': 'You',
     'name2': 'AI',
     'user_bio': '',

modules/ui.py

@@ -251,6 +251,7 @@ def list_interface_input_elements():
         'chat_style',
         'chat-instruct_command',
         'character_menu',
+        'user_menu',
         'name2',
         'context',
         'greeting',
@@ -353,6 +354,8 @@ def save_settings(state, preset, extensions_list, show_controls, theme_state, ma
     output['preset'] = preset
     output['prompt-notebook'] = state['prompt_menu-default'] if state['show_two_notebook_columns'] else state['prompt_menu-notebook']
     output['character'] = state['character_menu']
+    if 'user_menu' in state and state['user_menu']:
+        output['user'] = state['user_menu']
     output['seed'] = int(output['seed'])
     output['show_controls'] = show_controls
     output['dark_theme'] = True if theme_state == 'dark' else False
@@ -457,6 +460,7 @@ def setup_auto_save():
         'chat_style',
         'chat-instruct_command',
         'character_menu',
+        'user_menu',
         'name1',
         'name2',
         'context',

modules/ui_chat.py

@@ -137,6 +137,12 @@ def create_character_settings_ui():
                     shared.gradio['greeting'] = gr.Textbox(value=shared.settings['greeting'], lines=5, label='Greeting', elem_classes=['add_scrollbar'], elem_id="character-greeting")
 
                 with gr.Tab("User"):
+                    with gr.Row():
+                        shared.gradio['user_menu'] = gr.Dropdown(value=shared.settings['user'], choices=utils.get_available_users(), label='User', elem_id='user-menu', info='Select a user profile.', elem_classes='slim-dropdown')
+                        ui.create_refresh_button(shared.gradio['user_menu'], lambda: None, lambda: {'choices': utils.get_available_users()}, 'refresh-button', interactive=not mu)
+                        shared.gradio['save_user'] = gr.Button('💾', elem_classes='refresh-button', elem_id="save-user", interactive=not mu)
+                        shared.gradio['delete_user'] = gr.Button('🗑️', elem_classes='refresh-button', interactive=not mu)
+
                     shared.gradio['name1'] = gr.Textbox(value=shared.settings['name1'], lines=1, label='Name')
                     shared.gradio['user_bio'] = gr.Textbox(value=shared.settings['user_bio'], lines=10, label='Description', info='Here you can optionally write a description of yourself.', placeholder='{{user}}\'s personality: ...', elem_classes=['add_scrollbar'], elem_id="user-description")
 
@@ -372,3 +378,11 @@ def create_event_handlers():
         gradio('enable_web_search'),
         gradio('web_search_row')
     )
+
+    # User menu event handlers
+    shared.gradio['user_menu'].change(
+        ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(
+        chat.handle_user_menu_change, gradio('interface_state'), gradio('name1', 'user_bio', 'your_picture'), show_progress=False)
+
+    shared.gradio['save_user'].click(chat.handle_save_user_click, gradio('name1'), gradio('save_user_filename', 'user_saver'), show_progress=False)
+    shared.gradio['delete_user'].click(lambda: gr.update(visible=True), None, gradio('user_deleter'), show_progress=False)

modules/ui_file_saving.py

@@ -39,6 +39,19 @@ def create_ui():
             shared.gradio['delete_character_cancel'] = gr.Button('Cancel', elem_classes="small-button")
             shared.gradio['delete_character_confirm'] = gr.Button('Delete', elem_classes="small-button", variant='stop', interactive=not mu)
 
+    # User saver/deleter
+    with gr.Group(visible=False, elem_classes='file-saver') as shared.gradio['user_saver']:
+        shared.gradio['save_user_filename'] = gr.Textbox(lines=1, label='File name', info='The user profile will be saved to your user_data/users folder with this base filename.')
+        with gr.Row():
+            shared.gradio['save_user_cancel'] = gr.Button('Cancel', elem_classes="small-button")
+            shared.gradio['save_user_confirm'] = gr.Button('Save', elem_classes="small-button", variant='primary', interactive=not mu)
+
+    with gr.Group(visible=False, elem_classes='file-saver') as shared.gradio['user_deleter']:
+        gr.Markdown('Confirm the user deletion?')
+        with gr.Row():
+            shared.gradio['delete_user_cancel'] = gr.Button('Cancel', elem_classes="small-button")
+            shared.gradio['delete_user_confirm'] = gr.Button('Delete', elem_classes="small-button", variant='stop', interactive=not mu)
+
     # Preset saver
     with gr.Group(visible=False, elem_classes='file-saver') as shared.gradio['preset_saver']:
         shared.gradio['save_preset_filename'] = gr.Textbox(lines=1, label='File name', info='The preset will be saved to your user_data/presets folder with this base filename.')
@@ -69,6 +82,12 @@ def create_event_handlers():
     shared.gradio['save_character_cancel'].click(lambda: gr.update(visible=False), None, gradio('character_saver'), show_progress=False)
     shared.gradio['delete_character_cancel'].click(lambda: gr.update(visible=False), None, gradio('character_deleter'), show_progress=False)
 
+    # User save/delete event handlers
+    shared.gradio['save_user_confirm'].click(handle_save_user_confirm_click, gradio('name1', 'user_bio', 'your_picture', 'save_user_filename'), gradio('user_menu', 'user_saver'), show_progress=False)
+    shared.gradio['delete_user_confirm'].click(handle_delete_user_confirm_click, gradio('user_menu'), gradio('user_menu', 'user_deleter'), show_progress=False)
+    shared.gradio['save_user_cancel'].click(lambda: gr.update(visible=False), None, gradio('user_saver'), show_progress=False)
+    shared.gradio['delete_user_cancel'].click(lambda: gr.update(visible=False), None, gradio('user_deleter'), show_progress=False)
+
 
 def handle_save_preset_confirm_click(filename, contents):
     try:
@@ -165,3 +184,33 @@ def handle_delete_grammar_click(grammar_file):
         "user_data/grammars/",
         gr.update(visible=True)
     ]
+
+
+def handle_save_user_confirm_click(name1, user_bio, your_picture, filename):
+    try:
+        chat.save_user(name1, user_bio, your_picture, filename)
+        available_users = utils.get_available_users()
+        output = gr.update(choices=available_users, value=filename)
+    except Exception:
+        output = gr.update()
+        traceback.print_exc()
+
+    return [
+        output,
+        gr.update(visible=False)
+    ]
+
+
+def handle_delete_user_confirm_click(user):
+    try:
+        index = str(utils.get_available_users().index(user))
+        chat.delete_user(user)
+        output = chat.update_user_menu_after_deletion(index)
+    except Exception:
+        output = gr.update()
+        traceback.print_exc()
+
+    return [
+        output,
+        gr.update(visible=False)
+    ]

modules/utils.py

@@ -219,6 +219,13 @@ def get_available_characters():
     return sorted(set((k.stem for k in paths)), key=natural_keys)
 
 
+def get_available_users():
+    users_dir = Path('user_data/users')
+    users_dir.mkdir(parents=True, exist_ok=True)
+    paths = (x for x in users_dir.iterdir() if x.suffix in ('.json', '.yaml', '.yml'))
+    return sorted(set((k.stem for k in paths)), key=natural_keys)
+
+
 def get_available_instruction_templates():
     path = "user_data/instruction-templates"
     paths = []

user_data/users/Default.yaml

@@ -0,0 +1,2 @@
+name: You
+user_bio: ''