From 49c60882bf8f4806d0f47e890d7f434007941796 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Wed, 3 Dec 2025 16:07:51 -0800
Subject: [PATCH] Image generation: Safer image uploading

---
 modules/chat.py                | 16 +---------------
 modules/image_utils.py         | 20 ++++++++++++++++----
 modules/ui_image_generation.py |  4 ++--
 3 files changed, 19 insertions(+), 21 deletions(-)

diff --git a/modules/chat.py b/modules/chat.py
index 55984d7a..acfc2f66 100644
--- a/modules/chat.py
+++ b/modules/chat.py
@@ -3,7 +3,6 @@ import copy
 import functools
 import html
 import json
-import os
 import pprint
 import re
 import shutil
@@ -26,6 +25,7 @@ from modules.html_generator import (
     convert_to_markdown,
     make_thumbnail
 )
+from modules.image_utils import open_image_safely
 from modules.logging_colors import logger
 from modules.text_generation import (
     generate_reply,
@@ -1516,20 +1516,6 @@ def load_instruction_template_memoized(template):
     return load_instruction_template(template)
 
 
-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')
diff --git a/modules/image_utils.py b/modules/image_utils.py
index 658f00d7..d2809fef 100644
--- a/modules/image_utils.py
+++ b/modules/image_utils.py
@@ -1,9 +1,7 @@
-"""
-Shared image processing utilities for multimodal support.
-Used by both ExLlamaV3 and llama.cpp implementations.
-"""
 import base64
 import io
+import os
+from pathlib import Path
 from typing import Any, List, Tuple
 
 from PIL import Image
@@ -11,6 +9,20 @@ from PIL import Image
 from modules.logging_colors import logger
 
 
+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 convert_pil_to_base64(image: Image.Image) -> str:
     """Converts a PIL Image to a base64 encoded string."""
     buffered = io.BytesIO()
diff --git a/modules/ui_image_generation.py b/modules/ui_image_generation.py
index 3a64bc8b..06ef3e82 100644
--- a/modules/ui_image_generation.py
+++ b/modules/ui_image_generation.py
@@ -7,7 +7,6 @@ from pathlib import Path
 
 import gradio as gr
 import numpy as np
-from PIL import Image
 from PIL.PngImagePlugin import PngInfo
 
 from modules import shared, ui, utils
@@ -16,6 +15,7 @@ from modules.image_models import (
     load_image_model,
     unload_image_model
 )
+from modules.image_utils import open_image_safely
 from modules.logging_colors import logger
 from modules.text_generation import stop_everything_event
 from modules.utils import gradio
@@ -159,7 +159,7 @@ def save_generated_images(images, state, actual_seed):
 def read_image_metadata(image_path):
     """Read generation metadata from PNG file."""
     try:
-        with Image.open(image_path) as img:
+        with open_image_safely(image_path) as img:
             if hasattr(img, 'text') and 'image_gen_settings' in img.text:
                 return json.loads(img.text['image_gen_settings'])
     except Exception as e: