Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion css/main.css
Original file line number Diff line number Diff line change
Expand Up @@ -1645,7 +1645,7 @@ button:focus {
}

#user-description textarea {
height: calc(100vh - 231px) !important;
height: calc(100vh - 334px) !important;
min-height: 90px !important;
}

Expand Down
151 changes: 150 additions & 1 deletion modules/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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) -%}
Expand Down
1 change: 1 addition & 0 deletions modules/shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -298,6 +298,7 @@

# Character settings
'character': 'Assistant',
'user': 'Default',
'name1': 'You',
'name2': 'AI',
'user_bio': '',
Expand Down
4 changes: 4 additions & 0 deletions modules/ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -251,6 +251,7 @@ def list_interface_input_elements():
'chat_style',
'chat-instruct_command',
'character_menu',
'user_menu',
'name2',
'context',
'greeting',
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -457,6 +460,7 @@ def setup_auto_save():
'chat_style',
'chat-instruct_command',
'character_menu',
'user_menu',
'name1',
'name2',
'context',
Expand Down
14 changes: 14 additions & 0 deletions modules/ui_chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down Expand Up @@ -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)
49 changes: 49 additions & 0 deletions modules/ui_file_saving.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.')
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
]
7 changes: 7 additions & 0 deletions modules/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []
Expand Down
2 changes: 2 additions & 0 deletions user_data/users/Default.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
name: You
user_bio: ''