diff --git a/modules/chat.py b/modules/chat.py index e4fcaabe..e37c7a4e 100644 --- a/modules/chat.py +++ b/modules/chat.py @@ -2612,19 +2612,23 @@ def handle_load_template_click(instruction_template): def handle_save_template_click(instruction_template_str): import gradio as gr contents = generate_instruction_template_yaml(instruction_template_str) + root = str(shared.user_data_dir / 'instruction-templates') + '/' return [ "My Template.yaml", - str(shared.user_data_dir / 'instruction-templates') + '/', + root, contents, + root, gr.update(visible=True) ] def handle_delete_template_click(template): import gradio as gr + root = str(shared.user_data_dir / 'instruction-templates') + '/' return [ f"{template}.yaml", - str(shared.user_data_dir / 'instruction-templates') + '/', + root, + root, gr.update(visible=False) ] diff --git a/modules/image_utils.py b/modules/image_utils.py index d2809fef..b3138790 100644 --- a/modules/image_utils.py +++ b/modules/image_utils.py @@ -77,7 +77,18 @@ def process_message_content(content: Any) -> Tuple[str, List[Image.Image]]: # Support external URLs try: import requests - response = requests.get(image_url, timeout=10) + from urllib.parse import urljoin + from modules.web_search import _validate_url + _validate_url(image_url) + url = image_url + for _ in range(5): + response = requests.get(url, timeout=10, allow_redirects=False) + if response.is_redirect and 'Location' in response.headers: + url = urljoin(url, response.headers['Location']) + _validate_url(url) + else: + break + response.raise_for_status() image_data = response.content image = Image.open(io.BytesIO(image_data)) diff --git a/modules/ui_chat.py b/modules/ui_chat.py index d2a515b8..f1dc7883 100644 --- a/modules/ui_chat.py +++ b/modules/ui_chat.py @@ -350,13 +350,13 @@ def create_event_handlers(): shared.gradio['load_template'].click(chat.handle_load_template_click, gradio('instruction_template'), gradio('instruction_template_str', 'instruction_template'), show_progress=False) shared.gradio['save_template'].click( ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then( - chat.handle_save_template_click, gradio('instruction_template_str'), gradio('save_filename', 'save_root', 'save_contents', 'file_saver'), show_progress=False) + chat.handle_save_template_click, gradio('instruction_template_str'), gradio('save_filename', 'save_root', 'save_contents', 'save_root_state', 'file_saver'), show_progress=False) shared.gradio['restore_character'].click( ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then( chat.restore_character_for_ui, gradio('interface_state'), gradio('interface_state', 'name2', 'context', 'greeting', 'character_picture'), show_progress=False) - shared.gradio['delete_template'].click(chat.handle_delete_template_click, gradio('instruction_template'), gradio('delete_filename', 'delete_root', 'file_deleter'), show_progress=False) + shared.gradio['delete_template'].click(chat.handle_delete_template_click, gradio('instruction_template'), gradio('delete_filename', 'delete_root', 'delete_root_state', 'file_deleter'), show_progress=False) shared.gradio['save_chat_history'].click( lambda x: json.dumps(x, indent=4), gradio('history'), gradio('temporary_text')).then( None, gradio('temporary_text', 'character_menu', 'mode'), None, js=f'(hist, char, mode) => {{{ui.save_files_js}; saveHistory(hist, char, mode)}}') diff --git a/modules/ui_file_saving.py b/modules/ui_file_saving.py index 3ed256f8..99c4edd5 100644 --- a/modules/ui_file_saving.py +++ b/modules/ui_file_saving.py @@ -9,6 +9,12 @@ from modules.utils import gradio, sanitize_filename def create_ui(): mu = shared.args.multi_user + # Server-side per-session root paths for the generic file saver/deleter. + # Set by the handler that opens the dialog, read by the confirm handler. + # Using gr.State so they are session-scoped and safe for multi-user. + shared.gradio['save_root_state'] = gr.State(None) + shared.gradio['delete_root_state'] = gr.State(None) + # Text file saver with gr.Group(visible=False, elem_classes='file-saver') as shared.gradio['file_saver']: shared.gradio['save_filename'] = gr.Textbox(lines=1, label='File name') @@ -66,13 +72,13 @@ def create_event_handlers(): ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then( handle_save_preset_click, gradio('interface_state'), gradio('save_preset_contents', 'save_preset_filename', 'preset_saver'), show_progress=False) - shared.gradio['delete_preset'].click(handle_delete_preset_click, gradio('preset_menu'), gradio('delete_filename', 'delete_root', 'file_deleter'), show_progress=False) - shared.gradio['save_grammar'].click(handle_save_grammar_click, gradio('grammar_string'), gradio('save_contents', 'save_filename', 'save_root', 'file_saver'), show_progress=False) - shared.gradio['delete_grammar'].click(handle_delete_grammar_click, gradio('grammar_file'), gradio('delete_filename', 'delete_root', 'file_deleter'), show_progress=False) + shared.gradio['delete_preset'].click(handle_delete_preset_click, gradio('preset_menu'), gradio('delete_filename', 'delete_root', 'delete_root_state', 'file_deleter'), show_progress=False) + shared.gradio['save_grammar'].click(handle_save_grammar_click, gradio('grammar_string'), gradio('save_contents', 'save_filename', 'save_root', 'save_root_state', 'file_saver'), show_progress=False) + shared.gradio['delete_grammar'].click(handle_delete_grammar_click, gradio('grammar_file'), gradio('delete_filename', 'delete_root', 'delete_root_state', 'file_deleter'), show_progress=False) shared.gradio['save_preset_confirm'].click(handle_save_preset_confirm_click, gradio('save_preset_filename', 'save_preset_contents'), gradio('preset_menu', 'preset_saver'), show_progress=False) - shared.gradio['save_confirm'].click(handle_save_confirm_click, gradio('save_root', 'save_filename', 'save_contents'), gradio('file_saver'), show_progress=False) - shared.gradio['delete_confirm'].click(handle_delete_confirm_click, gradio('delete_root', 'delete_filename'), gradio('file_deleter'), show_progress=False) + shared.gradio['save_confirm'].click(handle_save_confirm_click, gradio('save_root_state', 'save_filename', 'save_contents'), gradio('save_root_state', 'file_saver'), show_progress=False) + shared.gradio['delete_confirm'].click(handle_delete_confirm_click, gradio('delete_root_state', 'delete_filename'), gradio('delete_root_state', 'file_deleter'), show_progress=False) shared.gradio['save_character_confirm'].click(handle_save_character_confirm_click, gradio('name2', 'greeting', 'context', 'character_picture', 'save_character_filename'), gradio('character_menu', 'character_saver'), show_progress=False) shared.gradio['delete_character_confirm'].click(handle_delete_character_confirm_click, gradio('character_menu'), gradio('character_menu', 'character_deleter'), show_progress=False) @@ -105,24 +111,30 @@ def handle_save_preset_confirm_click(filename, contents): ] -def handle_save_confirm_click(root, filename, contents): +def handle_save_confirm_click(root_state, filename, contents): try: + if root_state is None: + return None, gr.update(visible=False) + filename = sanitize_filename(filename) - utils.save_file(root + filename, contents) + utils.save_file(root_state + filename, contents) except Exception: traceback.print_exc() - return gr.update(visible=False) + return None, gr.update(visible=False) -def handle_delete_confirm_click(root, filename): +def handle_delete_confirm_click(root_state, filename): try: + if root_state is None: + return None, gr.update(visible=False) + filename = sanitize_filename(filename) - utils.delete_file(root + filename) + utils.delete_file(root_state + filename) except Exception: traceback.print_exc() - return gr.update(visible=False) + return None, gr.update(visible=False) def handle_save_character_confirm_click(name2, greeting, context, character_picture, filename): @@ -165,26 +177,32 @@ def handle_save_preset_click(state): def handle_delete_preset_click(preset): + root = str(shared.user_data_dir / "presets") + "/" return [ f"{preset}.yaml", - str(shared.user_data_dir / "presets") + "/", + root, + root, gr.update(visible=True) ] def handle_save_grammar_click(grammar_string): + root = str(shared.user_data_dir / "grammars") + "/" return [ grammar_string, "My Fancy Grammar.gbnf", - str(shared.user_data_dir / "grammars") + "/", + root, + root, gr.update(visible=True) ] def handle_delete_grammar_click(grammar_file): + root = str(shared.user_data_dir / "grammars") + "/" return [ grammar_file, - str(shared.user_data_dir / "grammars") + "/", + root, + root, gr.update(visible=True) ] diff --git a/modules/ui_session.py b/modules/ui_session.py index e1807dea..897bfd28 100644 --- a/modules/ui_session.py +++ b/modules/ui_session.py @@ -30,7 +30,7 @@ def create_ui(): if not mu: shared.gradio['save_settings'].click( ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then( - handle_save_settings, gradio('interface_state', 'preset_menu', 'extensions_menu', 'show_controls', 'theme_state'), gradio('save_contents', 'save_filename', 'save_root', 'file_saver'), show_progress=False) + handle_save_settings, gradio('interface_state', 'preset_menu', 'extensions_menu', 'show_controls', 'theme_state'), gradio('save_contents', 'save_filename', 'save_root', 'save_root_state', 'file_saver'), show_progress=False) shared.gradio['toggle_dark_mode'].click( lambda x: 'dark' if x == 'light' else 'light', gradio('theme_state'), gradio('theme_state')).then( @@ -51,10 +51,12 @@ def create_ui(): def handle_save_settings(state, preset, extensions, show_controls, theme): contents = ui.save_settings(state, preset, extensions, show_controls, theme, manual_save=True) + root = str(shared.user_data_dir) + "/" return [ contents, "settings.yaml", - str(shared.user_data_dir) + "/", + root, + root, gr.update(visible=True) ] diff --git a/modules/utils.py b/modules/utils.py index a14f8b8f..ff32e974 100644 --- a/modules/utils.py +++ b/modules/utils.py @@ -47,6 +47,10 @@ def save_file(fname, contents): logger.error(f'Invalid file path: \"{fname}\"') return + if Path(abs_path_str).suffix.lower() not in ('.yaml', '.yml', '.json', '.txt', '.gbnf'): + logger.error(f'Refusing to save file with disallowed extension: \"{fname}\"') + return + with open(abs_path_str, 'w', encoding='utf-8') as f: f.write(contents)