diff --git a/css/main.css b/css/main.css index ecf8568f0b..c55a29dcf2 100644 --- a/css/main.css +++ b/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; } diff --git a/modules/chat.py b/modules/chat.py index d1474cfec0..42c0d46d5a 100644 --- a/modules/chat.py +++ b/modules/chat.py @@ -32,7 +32,12 @@ 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) -%} diff --git a/modules/shared.py b/modules/shared.py index 0a27f33dc0..88e4b182ad 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -298,6 +298,7 @@ # Character settings 'character': 'Assistant', + 'user': 'Default', 'name1': 'You', 'name2': 'AI', 'user_bio': '', diff --git a/modules/ui.py b/modules/ui.py index 919a5740b5..8a49186139 100644 --- a/modules/ui.py +++ b/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', diff --git a/modules/ui_chat.py b/modules/ui_chat.py index c342ce5b0c..e00ddf5c39 100644 --- a/modules/ui_chat.py +++ b/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) diff --git a/modules/ui_file_saving.py b/modules/ui_file_saving.py index d1f9379b5a..720bfdecf1 100644 --- a/modules/ui_file_saving.py +++ b/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) + ] diff --git a/modules/utils.py b/modules/utils.py index b478f0664e..d366784730 100644 --- a/modules/utils.py +++ b/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 = [] diff --git a/user_data/users/Default.yaml b/user_data/users/Default.yaml new file mode 100644 index 0000000000..5c9dbacc49 --- /dev/null +++ b/user_data/users/Default.yaml @@ -0,0 +1,2 @@ +name: You +user_bio: ''