diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py
index ed4348414..91bc3f2d4 100644
--- a/backend/open_webui/routers/images.py
+++ b/backend/open_webui/routers/images.py
@@ -17,6 +17,8 @@ from fastapi.responses import FileResponse
 from open_webui.config import CACHE_DIR
 from open_webui.constants import ERROR_MESSAGES
 from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS
+
+from open_webui.models.chats import Chats
 from open_webui.routers.files import upload_file_handler, get_file_content_by_id
 from open_webui.utils.auth import get_admin_user, get_verified_user
 from open_webui.utils.headers import include_user_info_headers
@@ -510,14 +512,29 @@ def upload_image(request, image_data, content_type, metadata, user):
         process=False,
         user=user,
     )
+
+    if file_item and file_item.id:
+        # If chat_id and message_id are provided in metadata, link the file to the chat message
+        chat_id = metadata.get("chat_id")
+        message_id = metadata.get("message_id")
+
+        if chat_id and message_id:
+            Chats.insert_chat_files(
+                chat_id=chat_id,
+                message_id=message_id,
+                file_ids=[file_item.id],
+                user_id=user.id,
+            )
+
     url = request.app.url_path_for("get_file_content_by_id", id=file_item.id)
-    return url
+    return file_item, url
 
 
 @router.post("/generations")
 async def image_generations(
     request: Request,
     form_data: CreateImageForm,
+    metadata: Optional[dict] = None,
     user=Depends(get_verified_user),
 ):
     # if IMAGE_SIZE = 'auto', default WidthxHeight to the 512x512 default
@@ -535,6 +552,9 @@ async def image_generations(
         size = form_data.size
 
     width, height = tuple(map(int, size.split("x")))
+
+    metadata = metadata or {}
+
     model = get_image_model(request)
 
     r = None
@@ -595,7 +615,9 @@ async def image_generations(
                 else:
                     image_data, content_type = get_image_data(image["b64_json"])
 
-                url = upload_image(request, image_data, content_type, data, user)
+                _, url = upload_image(
+                    request, image_data, content_type, {**data, **metadata}, user
+                )
                 images.append({"url": url})
             return images
 
@@ -645,7 +667,9 @@ async def image_generations(
                     image_data, content_type = get_image_data(
                         image["bytesBase64Encoded"]
                     )
-                    url = upload_image(request, image_data, content_type, data, user)
+                    _, url = upload_image(
+                        request, image_data, content_type, {**data, **metadata}, user
+                    )
                     images.append({"url": url})
             elif model.endswith(":generateContent"):
                 for image in res["candidates"]:
@@ -654,8 +678,12 @@ async def image_generations(
                             image_data, content_type = get_image_data(
                                 part["inlineData"]["data"]
                             )
-                            url = upload_image(
-                                request, image_data, content_type, data, user
+                            _, url = upload_image(
+                                request,
+                                image_data,
+                                content_type,
+                                {**data, **metadata},
+                                user,
                             )
                             images.append({"url": url})
 
@@ -705,11 +733,11 @@ async def image_generations(
                     }
 
                 image_data, content_type = get_image_data(image["url"], headers)
-                url = upload_image(
+                _, url = upload_image(
                     request,
                     image_data,
                     content_type,
-                    form_data.model_dump(exclude_none=True),
+                    {**form_data.model_dump(exclude_none=True), **metadata},
                     user,
                 )
                 images.append({"url": url})
@@ -752,11 +780,11 @@ async def image_generations(
 
             for image in res["images"]:
                 image_data, content_type = get_image_data(image)
-                url = upload_image(
+                _, url = upload_image(
                     request,
                     image_data,
                     content_type,
-                    {**data, "info": res["info"]},
+                    {**data, "info": res["info"], **metadata},
                     user,
                 )
                 images.append({"url": url})
@@ -783,10 +811,13 @@ class EditImageForm(BaseModel):
 async def image_edits(
     request: Request,
     form_data: EditImageForm,
+    metadata: Optional[dict] = None,
     user=Depends(get_verified_user),
 ):
     size = None
     width, height = None, None
+    metadata = metadata or {}
+
     if (
         request.app.state.config.IMAGE_EDIT_SIZE
         and "x" in request.app.state.config.IMAGE_EDIT_SIZE
@@ -904,7 +935,9 @@ async def image_edits(
                 else:
                     image_data, content_type = get_image_data(image["b64_json"])
 
-                url = upload_image(request, image_data, content_type, data, user)
+                _, url = upload_image(
+                    request, image_data, content_type, {**data, **metadata}, user
+                )
                 images.append({"url": url})
             return images
 
@@ -957,8 +990,12 @@ async def image_edits(
                         image_data, content_type = get_image_data(
                             part["inlineData"]["data"]
                         )
-                        url = upload_image(
-                            request, image_data, content_type, data, user
+                        _, url = upload_image(
+                            request,
+                            image_data,
+                            content_type,
+                            {**data, **metadata},
+                            user,
                         )
                         images.append({"url": url})
 
@@ -1035,11 +1072,11 @@ async def image_edits(
                     }
 
                 image_data, content_type = get_image_data(image_url, headers)
-                url = upload_image(
+                _, url = upload_image(
                     request,
                     image_data,
                     content_type,
-                    form_data.model_dump(exclude_none=True),
+                    {**form_data.model_dump(exclude_none=True), **metadata},
                     user,
                 )
                 images.append({"url": url})
diff --git a/backend/open_webui/utils/files.py b/backend/open_webui/utils/files.py
index 2221d1707..a37ecf31c 100644
--- a/backend/open_webui/utils/files.py
+++ b/backend/open_webui/utils/files.py
@@ -14,6 +14,8 @@ from typing import Optional
 from pathlib import Path
 
 from open_webui.storage.provider import Storage
+
+from open_webui.models.chats import Chats
 from open_webui.models.files import Files
 from open_webui.routers.files import upload_file_handler
 
@@ -65,13 +67,14 @@ def get_image_url_from_base64(request, base64_image_string, metadata, user):
         # Extract base64 image data from the line
         image_data, content_type = get_image_data(base64_image_string)
         if image_data is not None:
-            image_url = upload_image(
+            _, image_url = upload_image(
                 request,
                 image_data,
                 content_type,
                 metadata,
                 user,
             )
+
         return image_url
     return None
 
diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index 865159208..05f76b172 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -799,6 +799,10 @@ async def chat_image_generation_handler(
             images = await image_edits(
                 request=request,
                 form_data=EditImageForm(**{"prompt": prompt, "image": input_images}),
+                metadata={
+                    "chat_id": metadata.get("chat_id", None),
+                    "message_id": metadata.get("message_id", None),
+                },
                 user=user,
             )
 
@@ -824,7 +828,7 @@ async def chat_image_generation_handler(
                 }
             )
 
-            system_message_content = "<context>The requested image has been created and is now being shown to the user. Let them know that it has been generated.</context>"
+            system_message_content = "<context>The requested image has been edited and created and is now being shown to the user. Let them know that it has been generated.</context>"
         except Exception as e:
             log.debug(e)
 
@@ -883,6 +887,10 @@ async def chat_image_generation_handler(
             images = await image_generations(
                 request=request,
                 form_data=CreateImageForm(**{"prompt": prompt}),
+                metadata={
+                    "chat_id": metadata.get("chat_id", None),
+                    "message_id": metadata.get("message_id", None),
+                },
                 user=user,
             )
 
@@ -2737,7 +2745,17 @@ async def process_chat_response(
 
                                         if ENABLE_CHAT_RESPONSE_BASE64_IMAGE_URL_CONVERSION:
                                             value = convert_markdown_base64_images(
-                                                request, value, metadata, user
+                                                request,
+                                                value,
+                                                {
+                                                    "chat_id": metadata.get(
+                                                        "chat_id", None
+                                                    ),
+                                                    "message_id": metadata.get(
+                                                        "message_id", None
+                                                    ),
+                                                },
+                                                user,
                                             )
 
                                         content = f"{content}{value}"