enh: image generation toggle

backend/open_webui/utils/middleware.py

@@ -31,6 +31,9 @@ from open_webui.routers.tasks import (
     generate_chat_tags,
 )
 from open_webui.routers.retrieval import process_web_search, SearchForm
+from open_webui.routers.images import image_generations, GenerateImageForm
+
+
 from open_webui.utils.webhook import post_webhook
 
 
@@ -486,6 +489,67 @@ async def chat_web_search_handler(
     return form_data
 
 
+async def chat_image_generation_handler(
+    request: Request, form_data: dict, extra_params: dict, user
+):
+    __event_emitter__ = extra_params["__event_emitter__"]
+    await __event_emitter__(
+        {
+            "type": "status",
+            "data": {"description": "Generating an image", "done": False},
+        }
+    )
+
+    messages = form_data["messages"]
+    user_message = get_last_user_message(messages)
+
+    system_message_content = ""
+
+    try:
+        images = await image_generations(
+            request=request,
+            form_data=GenerateImageForm(**{"prompt": user_message}),
+            user=user,
+        )
+
+        await __event_emitter__(
+            {
+                "type": "status",
+                "data": {"description": "Generated an image", "done": True},
+            }
+        )
+
+        for image in images:
+            await __event_emitter__(
+                {
+                    "type": "message",
+                    "data": {"content": f"![Generated Image]({image['url']})"},
+                }
+            )
+
+        system_message_content = "<context>User is shown the generated image, tell the user that the image has been generated</context>"
+    except Exception as e:
+        log.exception(e)
+        await __event_emitter__(
+            {
+                "type": "status",
+                "data": {
+                    "description": f"An error occured while generating an image",
+                    "done": True,
+                },
+            }
+        )
+
+        system_message_content = "<context>Unable to generate an image, tell the user that an error occured</context>"
+
+    if system_message_content:
+        form_data["messages"] = add_or_update_system_message(
+            system_message_content, form_data["messages"]
+        )
+
+    return form_data
+
+
 async def chat_completion_files_handler(
     request: Request, body: dict, user: UserModel
 ) -> tuple[dict, dict[str, list]]:
@@ -640,6 +704,11 @@ async def process_chat_payload(request, form_data, metadata, user, model):
                 request, form_data, extra_params, user
             )
 
+        if "image_generation" in features and features["image_generation"]:
+            form_data = await chat_image_generation_handler(
+                request, form_data, extra_params, user
+            )
+
     try:
         form_data, flags = await chat_completion_filter_functions_handler(
             request, form_data, model, extra_params

src/lib/components/chat/Chat.svelte

@@ -111,6 +111,7 @@
 	$: selectedModelIds = atSelectedModel !== undefined ? [atSelectedModel.id] : selectedModels;
 
 	let selectedToolIds = [];
+	let imageGenerationEnabled = false;
 	let webSearchEnabled = false;
 
 	let chat = null;
@@ -1533,6 +1534,7 @@
 				files: (files?.length ?? 0) > 0 ? files : undefined,
 				tool_ids: selectedToolIds.length > 0 ? selectedToolIds : undefined,
 				features: {
+					image_generation: imageGenerationEnabled,
 					web_search: webSearchEnabled
 				},
 
@@ -1935,6 +1937,7 @@
 								bind:prompt
 								bind:autoScroll
 								bind:selectedToolIds
+								bind:imageGenerationEnabled
 								bind:webSearchEnabled
 								bind:atSelectedModel
 								transparentBackground={$settings?.backgroundImageUrl ?? false}
@@ -1985,6 +1988,7 @@
 								bind:prompt
 								bind:autoScroll
 								bind:selectedToolIds
+								bind:imageGenerationEnabled
 								bind:webSearchEnabled
 								bind:atSelectedModel
 								transparentBackground={$settings?.backgroundImageUrl ?? false}

src/lib/components/chat/MessageInput.svelte

@@ -62,12 +62,15 @@
 	export let files = [];
 
 	export let selectedToolIds = [];
+
+	export let imageGenerationEnabled = false;
 	export let webSearchEnabled = false;
 
 	$: onChange({
 		prompt,
 		files,
 		selectedToolIds,
+		imageGenerationEnabled,
 		webSearchEnabled
 	});
 
@@ -642,6 +645,7 @@
 								<div class=" flex">
 									<div class="ml-1 self-end mb-1.5 flex space-x-1">
 										<InputMenu
+											bind:imageGenerationEnabled
 											bind:webSearchEnabled
 											bind:selectedToolIds
 											{screenCaptureHandler}

src/lib/components/chat/MessageInput/InputMenu.svelte

@@ -14,6 +14,7 @@
 	import GlobeAltSolid from '$lib/components/icons/GlobeAltSolid.svelte';
 	import WrenchSolid from '$lib/components/icons/WrenchSolid.svelte';
 	import CameraSolid from '$lib/components/icons/CameraSolid.svelte';
+	import PhotoSolid from '$lib/components/icons/PhotoSolid.svelte';
 
 	const i18n = getContext('i18n');
 
@@ -24,11 +25,25 @@
 	export let selectedToolIds: string[] = [];
 
 	export let webSearchEnabled: boolean;
+	export let imageGenerationEnabled: boolean;
+
 	export let onClose: Function;
 
 	let tools = {};
 	let show = false;
 
+	let showImageGeneration = false;
+
+	$: showImageGeneration =
+		$config?.features?.enable_image_generation &&
+		($user.role === 'admin' || $user?.permissions?.features?.image_generation);
+
+	let showWebSearch = false;
+
+	$: showWebSearch =
+		$config?.features?.enable_web_search &&
+		($user.role === 'admin' || $user?.permissions?.features?.web_search);
+
 	$: if (show) {
 		init();
 	}
@@ -63,7 +78,7 @@
 
 	<div slot="content">
 		<DropdownMenu.Content
-			class="w-full max-w-[200px] rounded-xl px-1 py-1  border-gray-300/30 dark:border-gray-700/50 z-50 bg-white dark:bg-gray-850 dark:text-white shadow"
+			class="w-full max-w-[220px] rounded-xl px-1 py-1  border-gray-300/30 dark:border-gray-700/50 z-50 bg-white dark:bg-gray-850 dark:text-white shadow"
 			sideOffset={15}
 			alignOffset={-8}
 			side="top"
@@ -114,7 +129,23 @@
 				<hr class="border-black/5 dark:border-white/5 my-1" />
 			{/if}
 
-			{#if $config?.features?.enable_web_search && ($user.role === 'admin' || $user?.permissions?.features?.web_search)}
+			{#if showImageGeneration}
+				<button
+					class="flex w-full justify-between gap-2 items-center px-3 py-2 text-sm font-medium cursor-pointer rounded-xl"
+					on:click={() => {
+						imageGenerationEnabled = !imageGenerationEnabled;
+					}}
+				>
+					<div class="flex-1 flex items-center gap-2">
+						<PhotoSolid />
+						<div class=" line-clamp-1">{$i18n.t('Image')}</div>
+					</div>
+
+					<Switch state={imageGenerationEnabled} />
+				</button>
+			{/if}
+
+			{#if showWebSearch}
 				<button
 					class="flex w-full justify-between gap-2 items-center px-3 py-2 text-sm font-medium cursor-pointer rounded-xl"
 					on:click={() => {
@@ -128,7 +159,9 @@
 
 					<Switch state={webSearchEnabled} />
 				</button>
+			{/if}
 
+			{#if showImageGeneration || showWebSearch}
 				<hr class="border-black/5 dark:border-white/5 my-1" />
 			{/if}
 

src/lib/components/chat/Placeholder.svelte

@@ -34,6 +34,7 @@
 	export let files = [];
 
 	export let selectedToolIds = [];
+	export let imageGenerationEnabled = false;
 	export let webSearchEnabled = false;
 
 	let models = [];
@@ -194,6 +195,7 @@
 					bind:prompt
 					bind:autoScroll
 					bind:selectedToolIds
+					bind:imageGenerationEnabled
 					bind:webSearchEnabled
 					bind:atSelectedModel
 					{transparentBackground}

src/lib/components/icons/PhotoSolid.svelte

@@ -0,0 +1,11 @@
+<script lang="ts">
+	export let className = 'size-4';
+</script>
+
+<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="currentColor" class={className}>
+	<path
+		fill-rule="evenodd"
+		d="M1.5 6a2.25 2.25 0 0 1 2.25-2.25h16.5A2.25 2.25 0 0 1 22.5 6v12a2.25 2.25 0 0 1-2.25 2.25H3.75A2.25 2.25 0 0 1 1.5 18V6ZM3 16.06V18c0 .414.336.75.75.75h16.5A.75.75 0 0 0 21 18v-1.94l-2.69-2.689a1.5 1.5 0 0 0-2.12 0l-.88.879.97.97a.75.75 0 1 1-1.06 1.06l-5.16-5.159a1.5 1.5 0 0 0-2.12 0L3 16.061Zm10.125-7.81a1.125 1.125 0 1 1 2.25 0 1.125 1.125 0 0 1-2.25 0Z"
+		clip-rule="evenodd"
+	/>
+</svg>