From 282aa1918907fceec7f903d3dc2bc8492ce8e885 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 9 Oct 2025 09:26:35 -0700 Subject: [PATCH] Safer profile picture uploading --- modules/chat.py | 55 ++++++++++++++++++++++++++++++++++++---------- modules/ui_chat.py | 8 +++---- 2 files changed, 48 insertions(+), 15 deletions(-) diff --git a/modules/chat.py b/modules/chat.py index 70c55b74..a52d96af 100644 --- a/modules/chat.py +++ b/modules/chat.py @@ -3,8 +3,10 @@ import copy import functools import html import json +import os import pprint import re +import shutil import time from datetime import datetime from functools import partial @@ -1385,12 +1387,17 @@ def generate_pfp_cache(character): for path in [Path(f"user_data/characters/{character}.{extension}") for extension in ['png', 'jpg', 'jpeg']]: if path.exists(): original_img = Image.open(path) - original_img.save(Path(f'{cache_folder}/pfp_character.png'), format='PNG') + # Define file paths + pfp_path = Path(f'{cache_folder}/pfp_character.png') + thumb_path = Path(f'{cache_folder}/pfp_character_thumb.png') + # Save main picture and thumbnail + original_img.save(pfp_path, format='PNG') thumb = make_thumbnail(original_img) - thumb.save(Path(f'{cache_folder}/pfp_character_thumb.png'), format='PNG') + thumb.save(thumb_path, format='PNG') - return thumb + # Return the path to the thumbnail, not the in-memory PIL Image object. + return str(thumb_path) return None @@ -1507,7 +1514,22 @@ def load_instruction_template_memoized(template): return load_instruction_template(template) -def upload_character(file, img, tavern=False): +def open_image_safely(path): + if path is None or not isinstance(path, str) or not Path(path).exists(): + return None + + if os.path.islink(path): + return None + + try: + return Image.open(path) + except Exception as e: + logger.error(f"Failed to open image file: {path}. Reason: {e}") + return None + + +def upload_character(file, img_path, tavern=False): + img = open_image_safely(img_path) decoded_file = file if isinstance(file, str) else file.decode('utf-8') try: data = json.loads(decoded_file) @@ -1554,12 +1576,17 @@ def build_pygmalion_style_context(data): return context -def upload_tavern_character(img, _json): +def upload_tavern_character(img_path, _json): _json = {'char_name': _json['name'], 'char_persona': _json['description'], 'char_greeting': _json['first_mes'], 'example_dialogue': _json['mes_example'], 'world_scenario': _json['scenario']} - return upload_character(json.dumps(_json), img, tavern=True) + return upload_character(json.dumps(_json), img_path, tavern=True) -def check_tavern_character(img): +def check_tavern_character(img_path): + img = open_image_safely(img_path) + + if img is None: + return "Invalid or disallowed image file.", None, None, gr.update(interactive=False) + if "chara" not in img.info: return "Not a TavernAI card", None, None, gr.update(interactive=False) @@ -1571,7 +1598,8 @@ def check_tavern_character(img): return _json['name'], _json['description'], _json, gr.update(interactive=True) -def upload_your_profile_picture(img): +def upload_your_profile_picture(img_path): + img = open_image_safely(img_path) cache_folder = Path(shared.args.disk_cache_dir) if not cache_folder.exists(): cache_folder.mkdir() @@ -1614,15 +1642,19 @@ def save_character(name, greeting, context, picture, filename): save_file(filepath, data) path_to_img = Path(f'user_data/characters/{filename}.png') if picture is not None: - picture.save(path_to_img) + # Copy the image file from its source path to the character folder + shutil.copy(picture, path_to_img) logger.info(f'Saved {path_to_img}.') def delete_character(name, instruct=False): + # Check for character data files for extension in ["yml", "yaml", "json"]: delete_file(Path(f'user_data/characters/{name}.{extension}')) - delete_file(Path(f'user_data/characters/{name}.png')) + # Check for character image files + for extension in ["png", "jpg", "jpeg"]: + delete_file(Path(f'user_data/characters/{name}.{extension}')) def jinja_template_from_old_format(params, verbose=False): @@ -1974,8 +2006,9 @@ def handle_character_menu_change(state): ] -def handle_character_picture_change(picture): +def handle_character_picture_change(picture_path): """Update or clear cache when character picture changes""" + picture = open_image_safely(picture_path) cache_folder = Path(shared.args.disk_cache_dir) if not cache_folder.exists(): cache_folder.mkdir() diff --git a/modules/ui_chat.py b/modules/ui_chat.py index 7c388607..c342ce5b 100644 --- a/modules/ui_chat.py +++ b/modules/ui_chat.py @@ -152,14 +152,14 @@ def create_character_settings_ui(): with gr.Tab('YAML or JSON'): with gr.Row(): shared.gradio['upload_json'] = gr.File(type='binary', file_types=['.json', '.yaml'], label='JSON or YAML File', interactive=not mu) - shared.gradio['upload_img_bot'] = gr.Image(type='pil', label='Profile Picture (optional)', interactive=not mu) + shared.gradio['upload_img_bot'] = gr.Image(type='filepath', label='Profile Picture (optional)', interactive=not mu) shared.gradio['Submit character'] = gr.Button(value='Submit', interactive=False) with gr.Tab('TavernAI PNG'): with gr.Row(): with gr.Column(): - shared.gradio['upload_img_tavern'] = gr.Image(type='pil', label='TavernAI PNG File', elem_id='upload_img_tavern', interactive=not mu) + shared.gradio['upload_img_tavern'] = gr.Image(type='filepath', label='TavernAI PNG File', elem_id='upload_img_tavern', interactive=not mu) shared.gradio['tavern_json'] = gr.State() with gr.Column(): shared.gradio['tavern_name'] = gr.Textbox(value='', lines=1, label='Name', interactive=False) @@ -168,8 +168,8 @@ def create_character_settings_ui(): shared.gradio['Submit tavern character'] = gr.Button(value='Submit', interactive=False) with gr.Column(scale=1): - shared.gradio['character_picture'] = gr.Image(label='Character picture', type='pil', interactive=not mu) - shared.gradio['your_picture'] = gr.Image(label='Your picture', type='pil', value=Image.open(Path('user_data/cache/pfp_me.png')) if Path('user_data/cache/pfp_me.png').exists() else None, interactive=not mu) + shared.gradio['character_picture'] = gr.Image(label='Character picture', type='filepath', interactive=not mu) + shared.gradio['your_picture'] = gr.Image(label='Your picture', type='filepath', value=Image.open(Path('user_data/cache/pfp_me.png')) if Path('user_data/cache/pfp_me.png').exists() else None, interactive=not mu) def create_chat_settings_ui():