diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index bc626161db..b40b808f39 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -1426,6 +1426,8 @@ IMAGE_EDIT_MODEL = os.getenv('IMAGE_EDIT_MODEL', '') IMAGE_EDIT_SIZE = os.getenv('IMAGE_EDIT_SIZE', '') +ENABLE_OPENAI_IMAGE_EDIT_NORMALIZATION = os.getenv('ENABLE_OPENAI_IMAGE_EDIT_NORMALIZATION', 'true').lower() == 'true' + IMAGES_EDIT_OPENAI_API_BASE_URL = os.getenv('IMAGES_EDIT_OPENAI_API_BASE_URL', OPENAI_API_BASE_URL) IMAGES_EDIT_OPENAI_API_VERSION = os.getenv('IMAGES_EDIT_OPENAI_API_VERSION', '') diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index c24cc3899e..3db3984d78 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -16,8 +16,10 @@ from urllib.parse import quote, urlparse import aiohttp from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile from fastapi.responses import FileResponse +from PIL import Image, ImageOps from open_webui.config import ( CACHE_DIR, + ENABLE_OPENAI_IMAGE_EDIT_NORMALIZATION, IMAGE_AUTO_SIZE_MODELS_REGEX_PATTERN, IMAGE_URL_RESPONSE_MODELS_REGEX_PATTERN, ) @@ -53,6 +55,14 @@ IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True) router = APIRouter() +IMAGE_FILE_EXTENSIONS = { + 'image/jpeg': '.jpg', + 'image/jpg': '.jpg', + 'image/mpo': '.jpg', + 'image/png': '.png', + 'image/webp': '.webp', +} + IMAGE_CONFIG_KEYS = { 'ENABLE_IMAGE_GENERATION': 'image_generation.enable', 'ENABLE_IMAGE_PROMPT_GENERATION': 'image_generation.prompt.enable', @@ -104,6 +114,61 @@ def config_updates(data: dict, key_map: dict[str, str]) -> dict: return {key_map[field]: value for field, value in data.items() if field in key_map} +def normalize_openai_edit_image_data_url(data_url: str) -> str: + if not data_url.startswith('data:') or ',' not in data_url: + return data_url + + header, encoded = data_url.split(',', 1) + mime_type = header.split(';')[0].lstrip('data:').lower() + if mime_type not in {'image/jpeg', 'image/jpg', 'image/mpo'}: + return data_url + + try: + image_bytes = base64.b64decode(encoded) + with Image.open(io.BytesIO(image_bytes)) as image: + orientation = image.getexif().get(274) + needs_normalization = ( + mime_type == 'image/mpo' + or image.format == 'MPO' + or getattr(image, 'n_frames', 1) > 1 + or orientation not in (None, 1) + or image.mode not in ('RGB', 'L') + ) + + if not needs_normalization: + return data_url + + image.seek(0) + image = ImageOps.exif_transpose(image) + if image.mode != 'RGB': + image = image.convert('RGB') + + output = io.BytesIO() + image.save(output, format='JPEG', quality=95) + normalized_image = base64.b64encode(output.getvalue()).decode('utf-8') + return f'data:image/jpeg;base64,{normalized_image}' + except Exception as e: + log.debug(f'Image edit normalization skipped: {e}') + + return data_url + + +def get_image_file_item(base64_string, param_name='image', normalize=False): + data = normalize_openai_edit_image_data_url(base64_string) if normalize else base64_string + header, encoded = data.split(',', 1) + mime_type = header.split(';')[0].lstrip('data:') or 'image/png' + image_data = base64.b64decode(encoded) + extension = IMAGE_FILE_EXTENSIONS.get(mime_type.lower()) or mimetypes.guess_extension(mime_type) or '.png' + return ( + param_name, + ( + f'{uuid.uuid4()}{extension}', + io.BytesIO(image_data), + mime_type, + ), + ) + + async def set_image_model(request: Request, model: str): log.info(f'Setting image model to {model}') await Config.upsert({'image_generation.model': model}) @@ -915,20 +980,6 @@ async def image_edits( detail=ERROR_MESSAGES.DEFAULT(e, 'Error loading image'), ) - def get_image_file_item(base64_string, param_name='image'): - data = base64_string - header, encoded = data.split(',', 1) - mime_type = header.split(';')[0].lstrip('data:') - image_data = base64.b64decode(encoded) - return ( - param_name, - ( - f'{uuid.uuid4()}.png', - io.BytesIO(image_data), - mime_type if mime_type else 'image/png', - ), - ) - try: if image_config.IMAGE_EDIT_ENGINE == 'openai': headers = { @@ -956,10 +1007,21 @@ async def image_edits( files = [] if isinstance(form_data.image, str): - files = [get_image_file_item(form_data.image)] + files = [ + get_image_file_item( + form_data.image, + normalize=ENABLE_OPENAI_IMAGE_EDIT_NORMALIZATION, + ) + ] elif isinstance(form_data.image, list): for img in form_data.image: - files.append(get_image_file_item(img, 'image[]')) + files.append( + get_image_file_item( + img, + 'image[]', + normalize=ENABLE_OPENAI_IMAGE_EDIT_NORMALIZATION, + ) + ) url_search_params = '' if image_config.IMAGES_EDIT_OPENAI_API_VERSION: